feat: add support for model variants across chat components and session management

This commit is contained in:
Bohdan Triapitsyn
2026-01-08 14:48:06 +02:00
parent 8ff7b61b8a
commit 8e777a94c2
11 changed files with 472 additions and 35 deletions
+87 -1
View File
@@ -24,6 +24,9 @@ interface ContextState {
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>;
@@ -42,8 +45,12 @@ interface ContextActions {
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;
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;
@@ -71,6 +78,7 @@ export const useContextStore = create<ContextStore>()(
sessionModelSelections: new Map(),
sessionAgentSelections: new Map(),
sessionAgentModelSelections: new Map(),
sessionAgentModelVariantSelections: new Map(),
currentAgentContext: new Map(),
sessionContextUsage: new Map(),
sessionAgentEditModes: new Map(),
@@ -129,8 +137,62 @@ export const useContextStore = create<ContextStore>()(
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}`);
},
analyzeAndSaveExternalSessionChoices: async (sessionId: string, agents: any[], messages: Map<string, { info: any; parts: any[] }[]>) => {
const { saveAgentModelForSession } = get();
const { saveAgentModelForSession, saveAgentModelVariantForSession } = get();
const agentLastChoices = new Map<
string,
@@ -212,6 +274,14 @@ export const useContextStore = create<ContextStore>()(
const agentName = extractAgentFromMessage(infoAny, assistantMessages.indexOf(message));
if (agentName && agents.find((a) => a.name === agentName)) {
const resolvedVariant = typeof infoAny.variant === 'string' && infoAny.variant.trim().length > 0
? infoAny.variant
: undefined;
if (resolvedVariant) {
saveAgentModelVariantForSession(sessionId, agentName, infoAny.providerID, infoAny.modelID, resolvedVariant);
}
const choice = {
providerId: infoAny.providerID,
modelId: infoAny.modelID,
@@ -474,6 +544,10 @@ export const useContextStore = create<ContextStore>()(
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())]),
@@ -487,6 +561,17 @@ export const useContextStore = create<ContextStore>()(
});
}
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[]]) => {
@@ -500,6 +585,7 @@ export const useContextStore = create<ContextStore>()(
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,
+8 -5
View File
@@ -20,9 +20,10 @@ import { useContextStore } from "./contextStore";
// Helper function to clean up pending user message metadata
const cleanupPendingUserMessageMeta = (
currentPending: Map<string, { mode?: string; providerID?: string; modelID?: string }>,
currentPending: Map<string, { mode?: string; providerID?: string; modelID?: string; variant?: string }>,
sessionId: string
): Map<string, { mode?: string; providerID?: string; modelID?: string }> => {
): Map<string, { mode?: string; providerID?: string; modelID?: string; variant?: string }> => {
const nextPending = new Map(currentPending);
nextPending.delete(sessionId);
return nextPending;
@@ -338,12 +339,12 @@ interface MessageState {
sessionCompactionUntil: Map<string, number>;
sessionAbortFlags: Map<string, SessionAbortRecord>;
pendingAssistantHeaderSessions: Set<string>;
pendingUserMessageMetaBySession: Map<string, { mode?: string; providerID?: string; modelID?: string }>;
pendingUserMessageMetaBySession: Map<string, { mode?: string; providerID?: string; modelID?: string; variant?: string }>;
}
interface MessageActions {
loadMessages: (sessionId: string) => Promise<void>;
sendMessage: (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>) => Promise<void>;
sendMessage: (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>, variant?: string) => Promise<void>;
abortCurrentOperation: (currentSessionId?: string) => Promise<void>;
_addStreamingPartImmediate: (sessionId: string, messageId: string, part: Part, role?: string, currentSessionId?: string) => void;
addStreamingPart: (sessionId: string, messageId: string, part: Part, role?: string, currentSessionId?: string) => void;
@@ -546,7 +547,7 @@ export const useMessageStore = create<MessageStore>()(
});
},
sendMessage: async (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>) => {
sendMessage: async (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>, variant?: string) => {
if (!currentSessionId) {
throw new Error("No session selected");
}
@@ -663,6 +664,7 @@ export const useMessageStore = create<MessageStore>()(
mode: typeof agent === 'string' && agent.trim().length > 0 ? agent.trim() : undefined,
providerID,
modelID,
variant: typeof variant === 'string' && variant.trim().length > 0 ? variant : undefined,
});
return { pendingAssistantHeaderSessions: next, pendingUserMessageMetaBySession: nextUserMeta };
});
@@ -684,6 +686,7 @@ export const useMessageStore = create<MessageStore>()(
modelID,
text: effectiveContent,
agent,
variant,
files: filePayloads.length > 0 ? filePayloads : undefined,
additionalParts: additionalPartsPayload,
agentMentions: agentMentionName ? [{ name: agentMentionName }] : undefined,
+5 -1
View File
@@ -131,7 +131,7 @@ export interface SessionStore {
unshareSession: (id: string) => Promise<Session | null>;
setCurrentSession: (id: string | null) => void;
loadMessages: (sessionId: string) => Promise<void>;
sendMessage: (content: string, providerID: string, modelID: string, agent?: string, attachments?: AttachedFile[], agentMentionName?: string, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>) => Promise<void>;
sendMessage: (content: string, providerID: string, modelID: string, agent?: string, attachments?: AttachedFile[], agentMentionName?: string, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>, variant?: string) => Promise<void>;
abortCurrentOperation: () => Promise<void>;
acknowledgeSessionAbort: (sessionId: string) => void;
armAbortPrompt: (durationMs?: number) => number | null;
@@ -172,8 +172,12 @@ export interface SessionStore {
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;
analyzeAndSaveExternalSessionChoices: (sessionId: string, agents: Array<{ name: string; [key: string]: unknown }>) => Promise<Map<string, { providerId: string; modelId: string; timestamp: number }>>;
isOpenChamberCreatedSession: (sessionId: string) => boolean;
markSessionAsOpenChamberCreated: (sessionId: string) => void;
+47 -2
View File
@@ -345,6 +345,7 @@ interface ConfigStore {
agents: Agent[];
currentProviderId: string;
currentModelId: string;
currentVariant: string | undefined;
currentAgentName: string | undefined;
selectedProviderId: string;
agentModelSelections: { [agentName: string]: { providerId: string; modelId: string } };
@@ -363,6 +364,9 @@ interface ConfigStore {
loadAgents: (options?: { directory?: string | null }) => Promise<boolean>;
setProvider: (providerId: string) => void;
setModel: (modelId: string) => void;
setCurrentVariant: (variant: string | undefined) => void;
cycleCurrentVariant: () => void;
getCurrentModelVariants: () => string[];
setAgent: (agentName: string | undefined) => void;
setSelectedProvider: (providerId: string) => void;
setSettingsDefaultModel: (model: string | undefined) => void;
@@ -399,6 +403,7 @@ export const useConfigStore = create<ConfigStore>()(
agents: [],
currentProviderId: "",
currentModelId: "",
currentVariant: undefined,
currentAgentName: undefined,
selectedProviderId: "",
agentModelSelections: {},
@@ -612,12 +617,12 @@ export const useConfigStore = create<ConfigStore>()(
agentModelSelections: state.agentModelSelections,
defaultProviders: state.defaultProviders,
};
const nextSnapshot: DirectoryScopedConfig = {
...baseSnapshot,
currentModelId: modelId,
};
return {
currentModelId: modelId,
directoryScoped: {
@@ -628,6 +633,46 @@ export const useConfigStore = create<ConfigStore>()(
});
},
setCurrentVariant: (variant: string | undefined) => {
set((state) => {
if (state.currentVariant === variant) {
return state;
}
return { currentVariant: variant };
});
},
getCurrentModelVariants: () => {
const model = get().getCurrentModel();
const variants = (model as { variants?: Record<string, unknown> } | undefined)?.variants;
if (!variants) {
return [];
}
return Object.keys(variants);
},
cycleCurrentVariant: () => {
const variantKeys = get().getCurrentModelVariants();
if (variantKeys.length === 0) {
return;
}
const current = get().currentVariant;
if (!current) {
set((state) => (state.currentVariant === variantKeys[0] ? state : { currentVariant: variantKeys[0] }));
return;
}
const index = variantKeys.indexOf(current);
if (index === -1 || index === variantKeys.length - 1) {
set((state) => (state.currentVariant === undefined ? state : { currentVariant: undefined }));
return;
}
const nextVariant = variantKeys[index + 1];
set((state) => (state.currentVariant === nextVariant ? state : { currentVariant: nextVariant }));
},
setSelectedProvider: (providerId: string) => {
set((state) => {
const directoryKey = state.activeDirectoryKey;
+35 -13
View File
@@ -290,7 +290,7 @@ export const useSessionStore = create<SessionStore>()(
get().evictLeastRecentlyUsed();
},
loadMessages: (sessionId: string) => useMessageStore.getState().loadMessages(sessionId),
sendMessage: async (content: string, providerID: string, modelID: string, agent?: string, attachments?: AttachedFile[], agentMentionName?: string, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>) => {
sendMessage: async (content: string, providerID: string, modelID: string, agent?: string, attachments?: AttachedFile[], agentMentionName?: string, additionalParts?: Array<{ text: string; attachments?: AttachedFile[] }>, variant?: string) => {
const draft = get().newSessionDraft;
const trimmedAgent = typeof agent === 'string' && agent.trim().length > 0 ? agent.trim() : undefined;
@@ -340,15 +340,25 @@ export const useSessionStore = create<SessionStore>()(
// ignored
}
if (draftProviderId && draftModelId) {
try {
useContextStore
.getState()
.saveAgentModelForSession(created.id, effectiveDraftAgent, draftProviderId, draftModelId);
} catch {
// ignored
if (draftProviderId && draftModelId) {
try {
useContextStore
.getState()
.saveAgentModelForSession(created.id, effectiveDraftAgent, draftProviderId, draftModelId);
} catch {
// ignored
}
if (variant !== undefined) {
try {
useContextStore
.getState()
.saveAgentModelVariantForSession(created.id, effectiveDraftAgent, draftProviderId, draftModelId, variant);
} catch {
// ignored
}
}
}
}
}
try {
@@ -365,7 +375,7 @@ export const useSessionStore = create<SessionStore>()(
try {
return await useMessageStore
.getState()
.sendMessage(content, providerID, modelID, effectiveDraftAgent, created.id, attachments, agentMentionName, additionalParts);
.sendMessage(content, providerID, modelID, effectiveDraftAgent, created.id, attachments, agentMentionName, additionalParts, variant);
} catch (error) {
setIdlePhase(created.id);
throw error;
@@ -385,14 +395,24 @@ export const useSessionStore = create<SessionStore>()(
} catch {
// ignored
}
}
if (variant !== undefined) {
try {
useContextStore
.getState()
.saveAgentModelVariantForSession(currentSessionId, effectiveAgent, providerID, modelID, variant);
} catch {
// ignored
}
}
}
if (currentSessionId) {
setBusyPhase(currentSessionId);
}
try {
return await useMessageStore.getState().sendMessage(content, providerID, modelID, effectiveAgent, currentSessionId || undefined, attachments, agentMentionName, additionalParts);
return await useMessageStore.getState().sendMessage(content, providerID, modelID, effectiveAgent, currentSessionId || undefined, attachments, agentMentionName, additionalParts, variant);
} catch (error) {
if (currentSessionId) {
setIdlePhase(currentSessionId);
@@ -473,7 +493,9 @@ export const useSessionStore = create<SessionStore>()(
getSessionAgentSelection: (sessionId: string) => useContextStore.getState().getSessionAgentSelection(sessionId),
saveAgentModelForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => useContextStore.getState().saveAgentModelForSession(sessionId, agentName, providerId, modelId),
getAgentModelForSession: (sessionId: string, agentName: string) => useContextStore.getState().getAgentModelForSession(sessionId, agentName),
analyzeAndSaveExternalSessionChoices: (sessionId: string, agents: Record<string, unknown>[]) => {
saveAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string, variant: string | undefined) => useContextStore.getState().saveAgentModelVariantForSession(sessionId, agentName, providerId, modelId, variant),
getAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => useContextStore.getState().getAgentModelVariantForSession(sessionId, agentName, providerId, modelId),
analyzeAndSaveExternalSessionChoices: (sessionId: string, agents: Record<string, unknown>[]) => {
const messages = useMessageStore.getState().messages;
return useContextStore.getState().analyzeAndSaveExternalSessionChoices(sessionId, agents, messages);
},