From 6e259c37db19417fb49db29acde2148ff97902de Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 2 Dec 2025 10:13:25 +0000 Subject: [PATCH] refactor: SessionProvider to use Zustand setters directly (Phase 2b) --- .../app/src/components/create-agent-modal.tsx | 9 +- .../app/src/contexts/realtime-context.tsx | 28 ++- packages/app/src/contexts/session-context.tsx | 211 +++++++++++++++--- 3 files changed, 195 insertions(+), 53 deletions(-) diff --git a/packages/app/src/components/create-agent-modal.tsx b/packages/app/src/components/create-agent-modal.tsx index bb52bc818..2ea0ea5e6 100644 --- a/packages/app/src/components/create-agent-modal.tsx +++ b/packages/app/src/components/create-agent-modal.tsx @@ -65,8 +65,9 @@ import type { WSInboundMessage, SessionOutboundMessage } from "@server/server/me import { formatConnectionStatus } from "@/utils/daemons"; import { trackAnalyticsEvent } from "@/utils/analytics"; import type { SessionContextValue } from "@/contexts/session-context"; -import { useSessionDirectory, useSessionForServer } from "@/hooks/use-session-directory"; +import { useSessionForServer } from "@/hooks/use-session-directory"; import type { UseWebSocketReturn } from "@/hooks/use-websocket"; +import { useSessionStore } from "@/stores/session-store"; export type CreateAgentInitialValues = { workingDir?: string; @@ -239,7 +240,7 @@ function AgentFlowModal({ const [selectedServerId, setSelectedServerId] = useState(initialServerId); const selectedSession = useSessionForServer(selectedServerId); const session = selectedSession ?? null; - const sessionDirectory = useSessionDirectory(); + const getSession = useSessionStore((state) => state.getSession); useEffect(() => { if (selectedServerId) { @@ -1484,7 +1485,7 @@ function AgentFlowModal({ clearQueuedProviderModelRequest(serverId); return; } - const daemonSession = sessionDirectory.get(serverId) ?? null; + const daemonSession = getSession(serverId) ?? null; if (!daemonSession?.ws?.isConnected) { return; } @@ -1500,7 +1501,7 @@ function AgentFlowModal({ queueProviderModelFetch, selectedProvider, selectedServerId, - sessionDirectory, + getSession, ]); const trimmedWorkingDir = workingDir.trim(); diff --git a/packages/app/src/contexts/realtime-context.tsx b/packages/app/src/contexts/realtime-context.tsx index 0ae489419..4acf77051 100644 --- a/packages/app/src/contexts/realtime-context.tsx +++ b/packages/app/src/contexts/realtime-context.tsx @@ -1,9 +1,9 @@ -import { createContext, useContext, useState, ReactNode, useCallback, useEffect, useRef, useMemo } from "react"; +import { createContext, useContext, useState, ReactNode, useCallback, useEffect, useRef } from "react"; import { useSpeechmaticsAudio } from "@/hooks/use-speechmatics-audio"; -import { useSessionDirectory } from "@/hooks/use-session-directory"; import type { SessionContextValue } from "./session-context"; import { generateMessageId } from "@/types/stream"; import type { WSInboundMessage } from "@server/server/messages"; +import { useSessionStore } from "@/stores/session-store"; interface RealtimeContextValue { isRealtimeMode: boolean; @@ -33,8 +33,19 @@ interface RealtimeProviderProps { } export function RealtimeProvider({ children }: RealtimeProviderProps) { - const sessionDirectory = useSessionDirectory(); + const getSession = useSessionStore((state) => state.getSession); const [activeServerId, setActiveServerId] = useState(null); + const activeSession = useSessionStore( + useCallback( + (state: ReturnType) => { + if (!activeServerId) { + return null; + } + return state.sessions[activeServerId] ?? null; + }, + [activeServerId] + ) + ); const realtimeSessionRef = useRef(null); const [isRealtimeMode, setIsRealtimeMode] = useState(false); const bargeInPlaybackStopRef = useRef(null); @@ -126,13 +137,6 @@ export function RealtimeProvider({ children }: RealtimeProviderProps) { session?.setVoiceDetectionFlags(realtimeAudio.isDetecting, realtimeAudio.isSpeaking); }, [realtimeAudio.isDetecting, realtimeAudio.isSpeaking]); - const activeSession = useMemo(() => { - if (!activeServerId) { - return null; - } - return sessionDirectory.get(activeServerId) ?? null; - }, [activeServerId, sessionDirectory]); - useEffect(() => { realtimeSessionRef.current = activeSession; }, [activeSession]); @@ -153,7 +157,7 @@ export function RealtimeProvider({ children }: RealtimeProviderProps) { const startRealtime = useCallback( async (serverId: string) => { - const session = sessionDirectory.get(serverId) ?? null; + const session = getSession(serverId) ?? null; if (!session) { throw new Error(`Host ${serverId} is not connected`); } @@ -179,7 +183,7 @@ export function RealtimeProvider({ children }: RealtimeProviderProps) { throw error; } }, - [realtimeAudio, sessionDirectory] + [getSession, realtimeAudio] ); const stopRealtime = useCallback(async () => { diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 133bdc574..3690f5e09 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -32,11 +32,8 @@ import { ScrollView } from "react-native"; import * as FileSystem from 'expo-file-system'; import { useDaemonConnections } from "./daemon-connections-context"; import { useSessionStore } from "@/stores/session-store"; -import { getNowMs, isPerfLoggingEnabled, measurePayload, perfLog } from "@/utils/perf"; import type { AgentDirectoryEntry } from "@/types/agent-directory"; -const SESSION_CONTEXT_LOG_TAG = "[SessionContext]"; - const derivePendingPermissionKey = (agentId: string, request: AgentPermissionRequest) => { const fallbackId = request.id || @@ -336,7 +333,7 @@ export interface SessionContextValue { messages: MessageEntry[]; setMessages: (messages: MessageEntry[] | ((prev: MessageEntry[]) => MessageEntry[])) => void; currentAssistantMessage: string; - setCurrentAssistantMessage: (message: string) => void; + setCurrentAssistantMessage: (message: string | ((prev: string) => string)) => void; agentStreamState: Map; setAgentStreamState: (state: Map | ((prev: Map) => Map)) => void; initializingAgents: Map; @@ -399,6 +396,47 @@ export interface SessionContextValue { respondToPermission: (agentId: string, requestId: string, response: AgentPermissionResponse) => void; } +type SyncSessionField = ( + key: K, + value: SessionContextValue[K] +) => void; + +function useSyncedSessionState( + key: K, + initialValue: SessionContextValue[K] | (() => SessionContextValue[K]), + syncField: SyncSessionField +): [ + SessionContextValue[K], + (valueOrUpdater: SessionContextValue[K] | ((prev: SessionContextValue[K]) => SessionContextValue[K])) => void +] { + const [state, setState] = useState(() => { + return typeof initialValue === "function" + ? (initialValue as () => SessionContextValue[K])() + : initialValue; + }); + + const setAndSync = useCallback( + (valueOrUpdater: SessionContextValue[K] | ((prev: SessionContextValue[K]) => SessionContextValue[K])) => { + setState((previous: SessionContextValue[K]) => { + const next = + typeof valueOrUpdater === "function" + ? (valueOrUpdater as (prev: SessionContextValue[K]) => SessionContextValue[K])(previous) + : valueOrUpdater; + + if (Object.is(previous, next)) { + return previous; + } + + syncField(key, next); + return next; + }); + }, + [key, syncField] + ); + + return [state, setAndSync]; +} + const SessionContext = createContext(null); interface SessionProviderProps { @@ -419,6 +457,30 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid const clearSession = useSessionStore((state) => state.clearSession); const setAgentDirectory = useSessionStore((state) => state.setAgentDirectory); const clearAgentDirectory = useSessionStore((state) => state.clearAgentDirectory); + const hasRegisteredSessionRef = useRef(false); + const pendingSessionUpdatesRef = useRef>({}); + const initialSessionValueRef = useRef(null); + + const syncSessionPartial = useCallback( + (partial: Partial) => { + if (!hasRegisteredSessionRef.current) { + pendingSessionUpdatesRef.current = { + ...pendingSessionUpdatesRef.current, + ...partial, + }; + return; + } + updateSession(serverId, partial); + }, + [serverId, updateSession] + ); + + const syncSessionField = useCallback( + (key, value) => { + syncSessionPartial({ [key]: value } as Partial); + }, + [syncSessionPartial] + ); useEffect(() => { if (ws.isConnected) { @@ -464,16 +526,19 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid isSpeaking: () => isSpeakingRef.current, }); - const [isPlayingAudio, setIsPlayingAudio] = useState(false); + const [isPlayingAudio, updateIsPlayingAudio] = useSyncedSessionState("isPlayingAudio", false, syncSessionField); + const setIsPlayingAudio = useCallback((playing: boolean) => { + updateIsPlayingAudio(playing); + }, [updateIsPlayingAudio]); const [focusedAgentOverride, setFocusedAgentOverride] = useState(null); const [orchestratorFocusedAgentId, setOrchestratorFocusedAgentId] = useState(null); - const [messages, setMessages] = useState([]); - const [currentAssistantMessage, setCurrentAssistantMessage] = useState(""); - const [agentStreamState, setAgentStreamState] = useState>(new Map()); - const [initializingAgents, setInitializingAgents] = useState>(new Map()); - const [hasHydratedAgents, setHasHydratedAgents] = useState(false); + const [messages, setMessages] = useSyncedSessionState("messages", () => [], syncSessionField); + const [currentAssistantMessage, setCurrentAssistantMessage] = useSyncedSessionState("currentAssistantMessage", "", syncSessionField); + const [agentStreamState, setAgentStreamState] = useSyncedSessionState("agentStreamState", () => new Map(), syncSessionField); + const [initializingAgents, setInitializingAgents] = useSyncedSessionState("initializingAgents", () => new Map(), syncSessionField); + const [hasHydratedAgents, setHasHydratedAgents] = useSyncedSessionState("hasHydratedAgents", false, syncSessionField); - const [agents, setAgents] = useState>(new Map()); + const [agents, setAgents] = useSyncedSessionState("agents", () => new Map(), syncSessionField); const lightweightAgentDirectory = useMemo( () => buildAgentDirectoryEntries(serverId, agents), [agents, serverId] @@ -481,11 +546,11 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid useEffect(() => { setAgentDirectory(serverId, lightweightAgentDirectory); }, [lightweightAgentDirectory, serverId, setAgentDirectory]); - const [commands, setCommands] = useState>(new Map()); - const [pendingPermissions, setPendingPermissions] = useState>(new Map()); - const [gitDiffs, setGitDiffs] = useState>(new Map()); - const [fileExplorer, setFileExplorer] = useState>(new Map()); - const [providerModels, setProviderModels] = useState>(new Map()); + const [commands, setCommands] = useSyncedSessionState("commands", () => new Map(), syncSessionField); + const [pendingPermissions, setPendingPermissions] = useSyncedSessionState("pendingPermissions", () => new Map(), syncSessionField); + const [gitDiffs, setGitDiffs] = useSyncedSessionState("gitDiffs", () => new Map(), syncSessionField); + const [fileExplorer, setFileExplorer] = useSyncedSessionState("fileExplorer", () => new Map(), syncSessionField); + const [providerModels, setProviderModels] = useSyncedSessionState("providerModels", () => new Map(), syncSessionField); const draftInputsRef = useRef>(new Map()); const getDraftInput = useCallback((agentId) => { return draftInputsRef.current.get(agentId); @@ -497,7 +562,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid images: draft.images, }); }, []); - const [queuedMessages, setQueuedMessages] = useState(new Map()); + const [queuedMessages, setQueuedMessages] = useSyncedSessionState("queuedMessages", () => new Map(), syncSessionField); const activeAudioGroupsRef = useRef>(new Set()); const previousAgentStatusRef = useRef>(new Map()); const providerModelRequestIdsRef = useRef>(new Map()); @@ -1893,34 +1958,106 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid ] ); - useEffect(() => { - setSession(serverId, { ...value }); - }, [serverId, setSession]); + if (initialSessionValueRef.current === null) { + initialSessionValueRef.current = value; + } useEffect(() => { - const payload = { ...value }; - const shouldLog = isPerfLoggingEnabled(); - const start = shouldLog ? getNowMs() : null; - updateSession(serverId, payload); - if (shouldLog && start !== null) { - const durationMs = getNowMs() - start; - const metrics = measurePayload(payload); - perfLog(SESSION_CONTEXT_LOG_TAG, { - event: "updateSession", - serverId, - durationMs: Number(durationMs.toFixed(2)), - payloadApproxBytes: metrics.approxBytes, - payloadFieldCount: metrics.fieldCount, - }); - } - }, [serverId, updateSession, value]); + syncSessionField("focusedAgentId", focusedAgentId); + }, [focusedAgentId, syncSessionField]); useEffect(() => { + syncSessionPartial({ + serverId, + ws, + audioPlayer, + }); + }, [audioPlayer, serverId, syncSessionPartial, ws]); + + const sessionFunctionBindings = useMemo( + () => ({ + setIsPlayingAudio, + setVoiceDetectionFlags, + setFocusedAgentId, + setMessages, + setCurrentAssistantMessage, + setAgentStreamState, + setAgents, + setCommands, + setPendingPermissions, + requestGitDiff, + requestDirectoryListing, + requestFilePreview, + navigateExplorerBack, + requestProviderModels, + restartServer, + initializeAgent, + refreshAgent, + cancelAgentRun, + deleteAgent, + sendAgentMessage, + sendAgentAudio, + createAgent, + resumeAgent, + setAgentMode, + respondToPermission, + getDraftInput, + saveDraftInput, + setQueuedMessages, + }), + [ + cancelAgentRun, + createAgent, + deleteAgent, + getDraftInput, + initializeAgent, + navigateExplorerBack, + refreshAgent, + requestDirectoryListing, + requestFilePreview, + requestGitDiff, + requestProviderModels, + respondToPermission, + resumeAgent, + restartServer, + saveDraftInput, + sendAgentAudio, + sendAgentMessage, + setAgentMode, + setAgentStreamState, + setAgents, + setCommands, + setCurrentAssistantMessage, + setFocusedAgentId, + setIsPlayingAudio, + setMessages, + setPendingPermissions, + setQueuedMessages, + setVoiceDetectionFlags, + ] + ); + + useEffect(() => { + syncSessionPartial(sessionFunctionBindings); + }, [sessionFunctionBindings, syncSessionPartial]); + + useEffect(() => { + const initialSnapshot = initialSessionValueRef.current as SessionContextValue; + const initialPayload = { + ...initialSnapshot, + ...pendingSessionUpdatesRef.current, + }; + pendingSessionUpdatesRef.current = {}; + setSession(serverId, initialPayload); + hasRegisteredSessionRef.current = true; + return () => { + hasRegisteredSessionRef.current = false; + pendingSessionUpdatesRef.current = {}; clearSession(serverId); clearAgentDirectory(serverId); }; - }, [clearAgentDirectory, clearSession, serverId]); + }, [clearAgentDirectory, clearSession, serverId, setSession]); return (