Tom RochetteandBohdan Triapitsyn
1505274f94
* perf(stores): defer safeStorage writes off the interaction path Session switches funnel every persisted store slice through safeStorage.setItem, and doing those large JSON.stringify writes synchronously blocked the main thread for over a second. Add a write-behind buffer that: - Defers each setItem/removeItem to a later task via setTimeout(0) so the click-to-paint path is not blocked. - Coalesces repeated writes to the same key into a single backing flush. - Serves pending values from memory so read-after-write stays consistent within the deferral window. - Flushes synchronously on pagehide/beforeunload/visibilitychange/freeze so deferred state survives tab close, reload, and the mobile freeze lifecycle. Adds a test covering write deferral, coalescing, and pending read serving. * fix(stores): defer persisted JSON serialization * fix(stores): defer direct safeStorage writes --------- Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>
512 lines
25 KiB
TypeScript
512 lines
25 KiB
TypeScript
/* eslint-disable @typescript-eslint/no-explicit-any */
|
|
import { create } from "zustand";
|
|
import { devtools, persist } from "zustand/middleware";
|
|
import type { EditPermissionMode } from "./types/sessionTypes";
|
|
import { getAgentDefaultEditPermission } from "./utils/permissionUtils";
|
|
import { extractTokensFromMessage } from "./utils/tokenUtils";
|
|
import { calculateContextUsage } from "./utils/contextUtils";
|
|
import { createDeferredSafeJSONStorage } from "./utils/safeStorage";
|
|
|
|
interface ContextUsage {
|
|
totalTokens: number;
|
|
percentage: number;
|
|
contextLimit: number;
|
|
outputLimit?: number;
|
|
normalizedOutput?: number;
|
|
thresholdLimit: number;
|
|
lastMessageId?: string;
|
|
}
|
|
|
|
interface ContextState {
|
|
|
|
sessionModelSelections: Map<string, { providerId: string; modelId: string }>;
|
|
sessionAgentSelections: Map<string, string>;
|
|
|
|
sessionAgentModelSelections: Map<string, Map<string, { providerId: string; modelId: string }>>;
|
|
|
|
// sessionId → agentName → "providerId/modelId" → variant
|
|
sessionAgentModelVariantSelections: Map<string, Map<string, Map<string, string>>>;
|
|
|
|
currentAgentContext: Map<string, string>;
|
|
|
|
sessionContextUsage: Map<string, ContextUsage>;
|
|
|
|
sessionAgentEditModes: Map<string, Map<string, EditPermissionMode>>;
|
|
hasHydrated: boolean;
|
|
}
|
|
|
|
interface ContextActions {
|
|
|
|
saveSessionModelSelection: (sessionId: string, providerId: string, modelId: string) => void;
|
|
getSessionModelSelection: (sessionId: string) => { providerId: string; modelId: string } | null;
|
|
saveSessionAgentSelection: (sessionId: string, agentName: string) => void;
|
|
getSessionAgentSelection: (sessionId: string) => string | null;
|
|
|
|
saveAgentModelForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => void;
|
|
getAgentModelForSession: (sessionId: string, agentName: string) => { providerId: string; modelId: string } | null;
|
|
|
|
saveAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string, variant: string | undefined) => void;
|
|
getAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => string | undefined;
|
|
|
|
|
|
getContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => ContextUsage | null;
|
|
|
|
updateSessionContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => void;
|
|
|
|
initializeSessionContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => void;
|
|
|
|
pollForTokenUpdates: (sessionId: string, messageId: string, messages: Map<string, { info: any; parts: any[] }[]>, maxAttempts?: number) => void;
|
|
|
|
getCurrentAgent: (sessionId: string) => string | undefined;
|
|
|
|
getSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => EditPermissionMode;
|
|
toggleSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => void;
|
|
setSessionAgentEditMode: (sessionId: string, agentName: string | undefined, mode: EditPermissionMode, defaultMode?: EditPermissionMode) => void;
|
|
}
|
|
|
|
type ContextStore = ContextState & ContextActions;
|
|
|
|
const EDIT_PERMISSION_SEQUENCE: EditPermissionMode[] = ['ask', 'allow', 'full'];
|
|
const GLOBAL_EDIT_MODE_SESSION_ID = '__global__';
|
|
|
|
export const useContextStore = create<ContextStore>()(
|
|
devtools(
|
|
persist(
|
|
(set, get) => ({
|
|
|
|
sessionModelSelections: new Map(),
|
|
sessionAgentSelections: new Map(),
|
|
sessionAgentModelSelections: new Map(),
|
|
sessionAgentModelVariantSelections: new Map(),
|
|
currentAgentContext: new Map(),
|
|
sessionContextUsage: new Map(),
|
|
sessionAgentEditModes: new Map(),
|
|
hasHydrated: typeof window === "undefined",
|
|
|
|
saveSessionModelSelection: (sessionId: string, providerId: string, modelId: string) => {
|
|
set((state) => {
|
|
const newSelections = new Map(state.sessionModelSelections);
|
|
newSelections.set(sessionId, { providerId, modelId });
|
|
return { sessionModelSelections: newSelections };
|
|
});
|
|
},
|
|
|
|
getSessionModelSelection: (sessionId: string) => {
|
|
const { sessionModelSelections } = get();
|
|
return sessionModelSelections.get(sessionId) || null;
|
|
},
|
|
|
|
saveSessionAgentSelection: (sessionId: string, agentName: string) => {
|
|
set((state) => {
|
|
const newSelections = new Map(state.sessionAgentSelections);
|
|
newSelections.set(sessionId, agentName);
|
|
|
|
// Keep a "current" agent context for components that only know sessionId.
|
|
// This is also used by external-session inference logic.
|
|
const nextAgentContext = new Map(state.currentAgentContext);
|
|
nextAgentContext.set(sessionId, agentName);
|
|
|
|
return {
|
|
sessionAgentSelections: newSelections,
|
|
currentAgentContext: nextAgentContext,
|
|
};
|
|
});
|
|
},
|
|
|
|
getSessionAgentSelection: (sessionId: string) => {
|
|
const { sessionAgentSelections } = get();
|
|
return sessionAgentSelections.get(sessionId) || null;
|
|
},
|
|
|
|
saveAgentModelForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => {
|
|
set((state) => {
|
|
const newSelections = new Map(state.sessionAgentModelSelections);
|
|
|
|
let agentMap = newSelections.get(sessionId);
|
|
if (!agentMap) {
|
|
agentMap = new Map();
|
|
} else {
|
|
|
|
agentMap = new Map(agentMap);
|
|
}
|
|
|
|
agentMap.set(agentName, { providerId, modelId });
|
|
|
|
newSelections.set(sessionId, agentMap);
|
|
|
|
return { sessionAgentModelSelections: newSelections };
|
|
});
|
|
},
|
|
|
|
getAgentModelForSession: (sessionId: string, agentName: string) => {
|
|
const { sessionAgentModelSelections } = get();
|
|
const agentMap = sessionAgentModelSelections.get(sessionId);
|
|
if (!agentMap) return null;
|
|
return agentMap.get(agentName) || null;
|
|
},
|
|
|
|
saveAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string, variant: string | undefined) => {
|
|
set((state) => {
|
|
const newSelections = new Map(state.sessionAgentModelVariantSelections);
|
|
|
|
let agentMap = newSelections.get(sessionId);
|
|
if (!agentMap) {
|
|
agentMap = new Map();
|
|
} else {
|
|
agentMap = new Map(agentMap);
|
|
}
|
|
|
|
let modelMap = agentMap.get(agentName);
|
|
if (!modelMap) {
|
|
modelMap = new Map();
|
|
} else {
|
|
modelMap = new Map(modelMap);
|
|
}
|
|
|
|
const modelKey = `${providerId}/${modelId}`;
|
|
|
|
if (variant === undefined) {
|
|
modelMap.delete(modelKey);
|
|
|
|
if (modelMap.size === 0) {
|
|
agentMap.delete(agentName);
|
|
|
|
if (agentMap.size === 0) {
|
|
newSelections.delete(sessionId);
|
|
} else {
|
|
newSelections.set(sessionId, agentMap);
|
|
}
|
|
} else {
|
|
agentMap.set(agentName, modelMap);
|
|
newSelections.set(sessionId, agentMap);
|
|
}
|
|
} else {
|
|
modelMap.set(modelKey, variant);
|
|
agentMap.set(agentName, modelMap);
|
|
newSelections.set(sessionId, agentMap);
|
|
}
|
|
|
|
return { sessionAgentModelVariantSelections: newSelections };
|
|
});
|
|
},
|
|
|
|
getAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => {
|
|
const { sessionAgentModelVariantSelections } = get();
|
|
const agentMap = sessionAgentModelVariantSelections.get(sessionId);
|
|
if (!agentMap) return undefined;
|
|
const modelMap = agentMap.get(agentName);
|
|
if (!modelMap) return undefined;
|
|
return modelMap.get(`${providerId}/${modelId}`);
|
|
},
|
|
|
|
getContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => {
|
|
if (!sessionId) return null;
|
|
|
|
const limitsUsage = calculateContextUsage(0, contextLimit, outputLimit);
|
|
const safeContext = limitsUsage.contextLimit;
|
|
const normalizedOutput = limitsUsage.normalizedOutput ?? 0;
|
|
const thresholdLimit = limitsUsage.thresholdLimit;
|
|
|
|
if (safeContext === 0 || thresholdLimit === 0) {
|
|
return null;
|
|
}
|
|
|
|
const sessionMessages = messages.get(sessionId) || [];
|
|
const assistantMessages = sessionMessages.filter(m => m.info.role === 'assistant');
|
|
|
|
if (assistantMessages.length === 0) return null;
|
|
|
|
const lastAssistantMessage = assistantMessages[assistantMessages.length - 1];
|
|
const lastMessageId = lastAssistantMessage.info.id;
|
|
|
|
const scheduleUsageUpdate = (usage: ContextUsage) => {
|
|
const runUpdate = () => {
|
|
set((state) => {
|
|
const existing = state.sessionContextUsage.get(sessionId) as ContextUsage | undefined;
|
|
if (
|
|
existing &&
|
|
existing.totalTokens === usage.totalTokens &&
|
|
existing.percentage === usage.percentage &&
|
|
existing.contextLimit === usage.contextLimit &&
|
|
(existing.outputLimit ?? 0) === (usage.outputLimit ?? 0) &&
|
|
(existing.normalizedOutput ?? 0) === (usage.normalizedOutput ?? 0) &&
|
|
existing.thresholdLimit === usage.thresholdLimit &&
|
|
existing.lastMessageId === usage.lastMessageId
|
|
) {
|
|
return state;
|
|
}
|
|
|
|
const newContextUsage = new Map(state.sessionContextUsage);
|
|
newContextUsage.set(sessionId, usage);
|
|
return { sessionContextUsage: newContextUsage };
|
|
});
|
|
};
|
|
|
|
if (typeof queueMicrotask === 'function') {
|
|
queueMicrotask(runUpdate);
|
|
} else if (typeof window !== 'undefined') {
|
|
window.setTimeout(runUpdate, 0);
|
|
} else {
|
|
setTimeout(runUpdate, 0);
|
|
}
|
|
};
|
|
|
|
const cachedUsage = get().sessionContextUsage.get(sessionId) as ContextUsage | undefined;
|
|
if (cachedUsage && cachedUsage.lastMessageId === lastMessageId) {
|
|
const cachedOutput = cachedUsage.normalizedOutput ?? cachedUsage.outputLimit ?? 0;
|
|
const limitsChanged =
|
|
cachedUsage.contextLimit !== safeContext ||
|
|
cachedOutput !== normalizedOutput ||
|
|
cachedUsage.thresholdLimit !== thresholdLimit;
|
|
|
|
if (limitsChanged && cachedUsage.totalTokens > 0) {
|
|
const newPercentage = (cachedUsage.totalTokens / thresholdLimit) * 100;
|
|
const recalculated: ContextUsage = {
|
|
totalTokens: cachedUsage.totalTokens,
|
|
percentage: Math.min(newPercentage, 100),
|
|
contextLimit: safeContext,
|
|
outputLimit: limitsUsage.outputLimit,
|
|
normalizedOutput,
|
|
thresholdLimit,
|
|
lastMessageId,
|
|
};
|
|
scheduleUsageUpdate(recalculated);
|
|
return recalculated;
|
|
}
|
|
|
|
if (!limitsChanged && cachedUsage.totalTokens > 0) {
|
|
return cachedUsage;
|
|
}
|
|
}
|
|
|
|
const totalTokens = extractTokensFromMessage(lastAssistantMessage);
|
|
|
|
if (totalTokens === 0) {
|
|
return cachedUsage || null;
|
|
}
|
|
|
|
const usage = calculateContextUsage(totalTokens, contextLimit, outputLimit);
|
|
const result: ContextUsage = {
|
|
totalTokens,
|
|
percentage: usage.percentage,
|
|
contextLimit: usage.contextLimit,
|
|
outputLimit: usage.outputLimit,
|
|
normalizedOutput: usage.normalizedOutput,
|
|
thresholdLimit: usage.thresholdLimit,
|
|
lastMessageId,
|
|
};
|
|
|
|
scheduleUsageUpdate(result);
|
|
|
|
return result;
|
|
},
|
|
|
|
updateSessionContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => {
|
|
const sessionMessages = messages.get(sessionId) || [];
|
|
const assistantMessages = sessionMessages.filter(m => m.info.role === 'assistant');
|
|
|
|
if (assistantMessages.length === 0) return;
|
|
|
|
const lastAssistantMessage = assistantMessages[assistantMessages.length - 1];
|
|
const totalTokens = extractTokensFromMessage(lastAssistantMessage);
|
|
|
|
if (totalTokens === 0) return;
|
|
|
|
const usage = calculateContextUsage(totalTokens, contextLimit, outputLimit);
|
|
|
|
set((state) => {
|
|
const newContextUsage = new Map(state.sessionContextUsage);
|
|
newContextUsage.set(sessionId, {
|
|
totalTokens,
|
|
percentage: usage.percentage,
|
|
contextLimit: usage.contextLimit,
|
|
outputLimit: usage.outputLimit,
|
|
normalizedOutput: usage.normalizedOutput,
|
|
thresholdLimit: usage.thresholdLimit,
|
|
lastMessageId: lastAssistantMessage.info.id,
|
|
});
|
|
return { sessionContextUsage: newContextUsage };
|
|
});
|
|
},
|
|
|
|
initializeSessionContextUsage: (sessionId: string, contextLimit: number, outputLimit: number, messages: Map<string, { info: any; parts: any[] }[]>) => {
|
|
const state = get();
|
|
const existingUsage = state.sessionContextUsage.get(sessionId);
|
|
|
|
if (!existingUsage || existingUsage.totalTokens === 0) {
|
|
get().updateSessionContextUsage(sessionId, contextLimit, outputLimit, messages);
|
|
}
|
|
},
|
|
|
|
pollForTokenUpdates: (sessionId: string, messageId: string, messages: Map<string, { info: any; parts: any[] }[]>, maxAttempts: number = 10) => {
|
|
let attempts = 0;
|
|
|
|
const poll = () => {
|
|
attempts++;
|
|
const sessionMessages = messages.get(sessionId) || [];
|
|
const message = sessionMessages.find(m => m.info.id === messageId);
|
|
|
|
if (message && message.info.role === 'assistant') {
|
|
const totalTokens = extractTokensFromMessage(message);
|
|
|
|
if (totalTokens > 0) {
|
|
|
|
get().updateSessionContextUsage(sessionId, 0, 0, messages);
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (attempts < maxAttempts) {
|
|
setTimeout(poll, 1000);
|
|
}
|
|
};
|
|
|
|
setTimeout(poll, 2000);
|
|
},
|
|
|
|
getCurrentAgent: (sessionId: string) => {
|
|
const { currentAgentContext } = get();
|
|
return currentAgentContext.get(sessionId);
|
|
},
|
|
|
|
getSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode: EditPermissionMode = getAgentDefaultEditPermission(agentName)) => {
|
|
if (!sessionId || !agentName) {
|
|
return defaultMode;
|
|
}
|
|
|
|
const sessionMap = get().sessionAgentEditModes.get(sessionId);
|
|
const override = sessionMap?.get(agentName);
|
|
if (override !== undefined) {
|
|
return override;
|
|
}
|
|
|
|
// Fallback: global (applies to all sessions)
|
|
if (sessionId !== GLOBAL_EDIT_MODE_SESSION_ID) {
|
|
const globalMap = get().sessionAgentEditModes.get(GLOBAL_EDIT_MODE_SESSION_ID);
|
|
const globalOverride = globalMap?.get(agentName);
|
|
if (globalOverride !== undefined) {
|
|
return globalOverride;
|
|
}
|
|
}
|
|
|
|
return defaultMode;
|
|
},
|
|
|
|
setSessionAgentEditMode: (sessionId: string, agentName: string | undefined, mode: EditPermissionMode, defaultMode: EditPermissionMode = getAgentDefaultEditPermission(agentName)) => {
|
|
if (!sessionId || !agentName) {
|
|
return;
|
|
}
|
|
|
|
const normalizedDefault: EditPermissionMode = defaultMode ?? 'ask';
|
|
if (normalizedDefault === 'deny' || mode === 'deny') {
|
|
return;
|
|
}
|
|
|
|
if (!EDIT_PERMISSION_SEQUENCE.includes(mode)) {
|
|
return;
|
|
}
|
|
|
|
set((state) => {
|
|
const nextMap = new Map(state.sessionAgentEditModes);
|
|
const agentMap = new Map(nextMap.get(sessionId) ?? new Map());
|
|
|
|
if (mode === normalizedDefault) {
|
|
agentMap.delete(agentName);
|
|
if (agentMap.size === 0) {
|
|
nextMap.delete(sessionId);
|
|
} else {
|
|
nextMap.set(sessionId, agentMap);
|
|
}
|
|
} else {
|
|
agentMap.set(agentName, mode);
|
|
nextMap.set(sessionId, agentMap);
|
|
}
|
|
|
|
return { sessionAgentEditModes: nextMap };
|
|
});
|
|
},
|
|
|
|
toggleSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode: EditPermissionMode = getAgentDefaultEditPermission(agentName)) => {
|
|
if (!sessionId || !agentName) {
|
|
return;
|
|
}
|
|
|
|
const normalizedDefault: EditPermissionMode = defaultMode ?? 'ask';
|
|
if (normalizedDefault === 'deny') {
|
|
return;
|
|
}
|
|
|
|
const currentMode = get().getSessionAgentEditMode(sessionId, agentName, normalizedDefault);
|
|
const currentIndex = EDIT_PERMISSION_SEQUENCE.indexOf(currentMode);
|
|
const fallbackIndex = EDIT_PERMISSION_SEQUENCE.indexOf(normalizedDefault);
|
|
const baseIndex = currentIndex >= 0 ? currentIndex : (fallbackIndex >= 0 ? fallbackIndex : 0);
|
|
const nextIndex = (baseIndex + 1) % EDIT_PERMISSION_SEQUENCE.length;
|
|
const nextMode = EDIT_PERMISSION_SEQUENCE[nextIndex];
|
|
|
|
get().setSessionAgentEditMode(sessionId, agentName, nextMode, normalizedDefault);
|
|
},
|
|
|
|
}),
|
|
{
|
|
name: "context-store",
|
|
storage: createDeferredSafeJSONStorage(),
|
|
partialize: (state) => ({
|
|
sessionModelSelections: Array.from(state.sessionModelSelections.entries()),
|
|
sessionAgentSelections: Array.from(state.sessionAgentSelections.entries()),
|
|
sessionAgentModelSelections: Array.from(state.sessionAgentModelSelections.entries()).map(([sessionId, agentMap]) => [sessionId, Array.from(agentMap.entries())]),
|
|
sessionAgentModelVariantSelections: Array.from(state.sessionAgentModelVariantSelections.entries()).map(([sessionId, agentMap]) => [
|
|
sessionId,
|
|
Array.from(agentMap.entries()).map(([agentName, modelMap]) => [agentName, Array.from(modelMap.entries())]),
|
|
]),
|
|
currentAgentContext: Array.from(state.currentAgentContext.entries()),
|
|
sessionContextUsage: Array.from(state.sessionContextUsage.entries()),
|
|
sessionAgentEditModes: Array.from(state.sessionAgentEditModes.entries()).map(([sessionId, agentMap]) => [sessionId, Array.from(agentMap.entries())]),
|
|
}),
|
|
merge: (persistedState: any, currentState) => {
|
|
|
|
const agentModelSelections = new Map();
|
|
if (persistedState?.sessionAgentModelSelections) {
|
|
persistedState.sessionAgentModelSelections.forEach(([sessionId, agentArray]: [string, any[]]) => {
|
|
agentModelSelections.set(sessionId, new Map(agentArray));
|
|
});
|
|
}
|
|
|
|
const agentModelVariantSelections = new Map();
|
|
if (persistedState?.sessionAgentModelVariantSelections) {
|
|
persistedState.sessionAgentModelVariantSelections.forEach(([sessionId, agentArray]: [string, any[]]) => {
|
|
const agentMap = new Map();
|
|
agentArray.forEach(([agentName, modelArray]: [string, any[]]) => {
|
|
agentMap.set(agentName, new Map(modelArray));
|
|
});
|
|
agentModelVariantSelections.set(sessionId, agentMap);
|
|
});
|
|
}
|
|
|
|
const agentEditModes = new Map();
|
|
if (persistedState?.sessionAgentEditModes) {
|
|
persistedState.sessionAgentEditModes.forEach(([sessionId, agentArray]: [string, any[]]) => {
|
|
agentEditModes.set(sessionId, new Map(agentArray));
|
|
});
|
|
}
|
|
|
|
return {
|
|
...currentState,
|
|
...(persistedState as object),
|
|
sessionModelSelections: new Map(persistedState?.sessionModelSelections || []),
|
|
sessionAgentSelections: new Map(persistedState?.sessionAgentSelections || []),
|
|
sessionAgentModelSelections: agentModelSelections,
|
|
sessionAgentModelVariantSelections: agentModelVariantSelections,
|
|
currentAgentContext: new Map(persistedState?.currentAgentContext || []),
|
|
sessionContextUsage: new Map(persistedState?.sessionContextUsage || []),
|
|
sessionAgentEditModes: agentEditModes,
|
|
hasHydrated: true,
|
|
};
|
|
},
|
|
}
|
|
),
|
|
{
|
|
name: "context-store",
|
|
}
|
|
)
|
|
);
|