516 lines
25 KiB
TypeScript
516 lines
25 KiB
TypeScript
/* eslint-disable @typescript-eslint/no-explicit-any */
|
|
import { create } from "zustand";
|
|
import { devtools, persist, createJSONStorage } 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 { getSafeStorage } 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 }>>;
|
|
|
|
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;
|
|
|
|
analyzeAndSaveExternalSessionChoices: (sessionId: string, agents: any[], messages: Map<string, { info: any; parts: any[] }[]>) => Promise<Map<string, { providerId: string; modelId: string; timestamp: number }>>;
|
|
|
|
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'];
|
|
|
|
export const useContextStore = create<ContextStore>()(
|
|
devtools(
|
|
persist(
|
|
(set, get) => ({
|
|
|
|
sessionModelSelections: new Map(),
|
|
sessionAgentSelections: new Map(),
|
|
sessionAgentModelSelections: 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);
|
|
return { sessionAgentSelections: newSelections };
|
|
});
|
|
},
|
|
|
|
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;
|
|
},
|
|
|
|
analyzeAndSaveExternalSessionChoices: async (sessionId: string, agents: any[], messages: Map<string, { info: any; parts: any[] }[]>) => {
|
|
const { saveAgentModelForSession } = get();
|
|
|
|
const agentLastChoices = new Map<
|
|
string,
|
|
{
|
|
providerId: string;
|
|
modelId: string;
|
|
timestamp: number;
|
|
}
|
|
>();
|
|
|
|
const extractAgentFromMessage = (messageInfo: any, messageIndex: number): string | null => {
|
|
|
|
if ("mode" in messageInfo && messageInfo.mode && typeof messageInfo.mode === "string") {
|
|
const modeAgent = agents.find((a) => a.name === messageInfo.mode);
|
|
if (modeAgent) {
|
|
return messageInfo.mode;
|
|
}
|
|
}
|
|
|
|
if (messageInfo.providerID && messageInfo.modelID) {
|
|
const matchingAgent = agents.find((agent) => agent.model?.providerID === messageInfo.providerID && agent.model?.modelID === messageInfo.modelID);
|
|
if (matchingAgent) {
|
|
return matchingAgent.name;
|
|
}
|
|
}
|
|
|
|
const { currentAgentContext } = get();
|
|
const contextAgent = currentAgentContext.get(sessionId);
|
|
if (contextAgent && agents.find((a) => a.name === contextAgent)) {
|
|
return contextAgent;
|
|
}
|
|
|
|
if (messageIndex > 0 && messageInfo.providerID && messageInfo.modelID) {
|
|
|
|
const sessionMessages = messages.get(sessionId) || [];
|
|
const assistantMessages = sessionMessages.filter((m) => m.info.role === "assistant").sort((a, b) => a.info.time.created - b.info.time.created);
|
|
|
|
for (let i = messageIndex - 1; i >= 0; i--) {
|
|
const prevMessage = assistantMessages[i];
|
|
const prevInfo = prevMessage.info as any;
|
|
if (prevInfo.providerID === messageInfo.providerID && prevInfo.modelID === messageInfo.modelID) {
|
|
|
|
if (prevInfo.mode && typeof prevInfo.mode === "string") {
|
|
const prevModeAgent = agents.find((a) => a.name === prevInfo.mode);
|
|
if (prevModeAgent) {
|
|
return prevInfo.mode;
|
|
}
|
|
}
|
|
|
|
const prevMatchingAgent = agents.find((agent) => agent.model?.providerID === prevInfo.providerID && agent.model?.modelID === prevInfo.modelID);
|
|
if (prevMatchingAgent) {
|
|
return prevMatchingAgent.name;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (messageInfo.providerID && messageInfo.modelID) {
|
|
const buildAgent = agents.find((a) => a.name === "build");
|
|
if (buildAgent) {
|
|
return "build";
|
|
}
|
|
}
|
|
|
|
return null;
|
|
};
|
|
|
|
const sessionMessages = messages.get(sessionId) || [];
|
|
|
|
const allMessages = sessionMessages.filter((m: any) => m.info.role === "assistant" || m.info.role === "user").sort((a: any, b: any) => a.info.time.created - b.info.time.created);
|
|
const assistantMessages = sessionMessages.filter((m: any) => m.info.role === "assistant").sort((a: any, b: any) => a.info.time.created - b.info.time.created);
|
|
|
|
for (let messageIndex = 0; messageIndex < allMessages.length; messageIndex++) {
|
|
const message = allMessages[messageIndex];
|
|
const { info } = message;
|
|
const infoAny = info as any;
|
|
|
|
if (infoAny.providerID && infoAny.modelID) {
|
|
const agentName = extractAgentFromMessage(infoAny, assistantMessages.indexOf(message));
|
|
|
|
if (agentName && agents.find((a) => a.name === agentName)) {
|
|
const choice = {
|
|
providerId: infoAny.providerID,
|
|
modelId: infoAny.modelID,
|
|
timestamp: info.time.created,
|
|
};
|
|
|
|
const existing = agentLastChoices.get(agentName);
|
|
if (!existing || choice.timestamp > existing.timestamp) {
|
|
agentLastChoices.set(agentName, choice);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (const [agentName, choice] of agentLastChoices) {
|
|
saveAgentModelForSession(sessionId, agentName, choice.providerId, choice.modelId);
|
|
}
|
|
|
|
return agentLastChoices;
|
|
},
|
|
|
|
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);
|
|
return override ?? 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: createJSONStorage(() => getSafeStorage()),
|
|
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())]),
|
|
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 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,
|
|
currentAgentContext: new Map(persistedState?.currentAgentContext || []),
|
|
sessionContextUsage: new Map(persistedState?.sessionContextUsage || []),
|
|
sessionAgentEditModes: agentEditModes,
|
|
hasHydrated: true,
|
|
};
|
|
},
|
|
}
|
|
),
|
|
{
|
|
name: "context-store",
|
|
}
|
|
)
|
|
);
|