refactor: SessionProvider to use Zustand setters directly (Phase 2b)

This commit is contained in:
Mohamed Boudra
2025-12-02 10:13:25 +00:00
parent fa1d7dd295
commit 6e259c37db
3 changed files with 195 additions and 53 deletions
@@ -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<string | null>(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();
+16 -12
View File
@@ -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<string | null>(null);
const activeSession = useSessionStore(
useCallback(
(state: ReturnType<typeof useSessionStore.getState>) => {
if (!activeServerId) {
return null;
}
return state.sessions[activeServerId] ?? null;
},
[activeServerId]
)
);
const realtimeSessionRef = useRef<SessionContextValue | null>(null);
const [isRealtimeMode, setIsRealtimeMode] = useState(false);
const bargeInPlaybackStopRef = useRef<number | null>(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 () => {
+174 -37
View File
@@ -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<string, StreamItem[]>;
setAgentStreamState: (state: Map<string, StreamItem[]> | ((prev: Map<string, StreamItem[]>) => Map<string, StreamItem[]>)) => void;
initializingAgents: Map<string, boolean>;
@@ -399,6 +396,47 @@ export interface SessionContextValue {
respondToPermission: (agentId: string, requestId: string, response: AgentPermissionResponse) => void;
}
type SyncSessionField = <K extends keyof SessionContextValue>(
key: K,
value: SessionContextValue[K]
) => void;
function useSyncedSessionState<K extends keyof SessionContextValue>(
key: K,
initialValue: SessionContextValue[K] | (() => SessionContextValue[K]),
syncField: SyncSessionField
): [
SessionContextValue[K],
(valueOrUpdater: SessionContextValue[K] | ((prev: SessionContextValue[K]) => SessionContextValue[K])) => void
] {
const [state, setState] = useState<SessionContextValue[K]>(() => {
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<SessionContextValue | null>(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<Partial<SessionContextValue>>({});
const initialSessionValueRef = useRef<SessionContextValue | null>(null);
const syncSessionPartial = useCallback(
(partial: Partial<SessionContextValue>) => {
if (!hasRegisteredSessionRef.current) {
pendingSessionUpdatesRef.current = {
...pendingSessionUpdatesRef.current,
...partial,
};
return;
}
updateSession(serverId, partial);
},
[serverId, updateSession]
);
const syncSessionField = useCallback<SyncSessionField>(
(key, value) => {
syncSessionPartial({ [key]: value } as Partial<SessionContextValue>);
},
[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<string | null>(null);
const [orchestratorFocusedAgentId, setOrchestratorFocusedAgentId] = useState<string | null>(null);
const [messages, setMessages] = useState<MessageEntry[]>([]);
const [currentAssistantMessage, setCurrentAssistantMessage] = useState("");
const [agentStreamState, setAgentStreamState] = useState<Map<string, StreamItem[]>>(new Map());
const [initializingAgents, setInitializingAgents] = useState<Map<string, boolean>>(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<string, StreamItem[]>(), syncSessionField);
const [initializingAgents, setInitializingAgents] = useSyncedSessionState("initializingAgents", () => new Map<string, boolean>(), syncSessionField);
const [hasHydratedAgents, setHasHydratedAgents] = useSyncedSessionState("hasHydratedAgents", false, syncSessionField);
const [agents, setAgents] = useState<Map<string, Agent>>(new Map());
const [agents, setAgents] = useSyncedSessionState("agents", () => new Map<string, Agent>(), 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<Map<string, Command>>(new Map());
const [pendingPermissions, setPendingPermissions] = useState<Map<string, PendingPermission>>(new Map());
const [gitDiffs, setGitDiffs] = useState<Map<string, string>>(new Map());
const [fileExplorer, setFileExplorer] = useState<Map<string, AgentFileExplorerState>>(new Map());
const [providerModels, setProviderModels] = useState<Map<AgentProvider, ProviderModelState>>(new Map());
const [commands, setCommands] = useSyncedSessionState("commands", () => new Map<string, Command>(), syncSessionField);
const [pendingPermissions, setPendingPermissions] = useSyncedSessionState("pendingPermissions", () => new Map<string, PendingPermission>(), syncSessionField);
const [gitDiffs, setGitDiffs] = useSyncedSessionState("gitDiffs", () => new Map<string, string>(), syncSessionField);
const [fileExplorer, setFileExplorer] = useSyncedSessionState("fileExplorer", () => new Map<string, AgentFileExplorerState>(), syncSessionField);
const [providerModels, setProviderModels] = useSyncedSessionState("providerModels", () => new Map<AgentProvider, ProviderModelState>(), syncSessionField);
const draftInputsRef = useRef<Map<string, DraftInput>>(new Map());
const getDraftInput = useCallback<SessionContextValue["getDraftInput"]>((agentId) => {
return draftInputsRef.current.get(agentId);
@@ -497,7 +562,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
images: draft.images,
});
}, []);
const [queuedMessages, setQueuedMessages] = useState<SessionContextValue["queuedMessages"]>(new Map());
const [queuedMessages, setQueuedMessages] = useSyncedSessionState("queuedMessages", () => new Map(), syncSessionField);
const activeAudioGroupsRef = useRef<Set<string>>(new Set());
const previousAgentStatusRef = useRef<Map<string, AgentLifecycleStatus>>(new Map());
const providerModelRequestIdsRef = useRef<Map<AgentProvider, string>>(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 (
<SessionContext.Provider value={value}>