Batch hydrate normalized messages for session load

packages/ui/src/stores/message-v2/session-info.ts

@@ -0,0 +1,139 @@
+import type { Provider } from "../../types/session"
+import { DEFAULT_MODEL_OUTPUT_LIMIT } from "../session-models"
+import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance } from "../session-state"
+import { messageStoreBus } from "./bus"
+import type { SessionUsageState } from "./types"
+
+function getLatestUsageEntry(usage?: SessionUsageState) {
+  if (!usage?.latestMessageId) return undefined
+  return usage.entries[usage.latestMessageId]
+}
+
+function resolveSelectedModel(instanceProviders: Provider[], providerId?: string, modelId?: string) {
+  if (!providerId || !modelId) return undefined
+  const provider = instanceProviders.find((p) => p.id === providerId)
+  return provider?.models.find((m) => m.id === modelId)
+}
+
+export function updateSessionInfo(instanceId: string, sessionId: string): void {
+  const instanceSessions = sessions().get(instanceId)
+  if (!instanceSessions) return
+  const session = instanceSessions.get(sessionId)
+  if (!session) return
+
+  const store = messageStoreBus.getOrCreate(instanceId)
+  const usage = store.getSessionUsage(sessionId)
+  const hasUsageEntries = Boolean(usage && Object.keys(usage.entries).length > 0)
+
+  let totalInputTokens = usage?.totalInputTokens ?? 0
+  let totalOutputTokens = usage?.totalOutputTokens ?? 0
+  let totalReasoningTokens = usage?.totalReasoningTokens ?? 0
+  let totalCost = usage?.totalCost ?? 0
+  let actualUsageTokens = usage?.actualUsageTokens ?? 0
+
+  const latestEntry = getLatestUsageEntry(usage)
+  let latestHasContextUsage = latestEntry?.hasContextUsage ?? false
+
+  const previousInfo = sessionInfoByInstance().get(instanceId)?.get(sessionId)
+  let contextWindow = 0
+  let contextAvailableTokens: number | null = null
+  let contextAvailableFromPrevious = false
+  let isSubscriptionModel = false
+
+  if (!hasUsageEntries && previousInfo) {
+    totalInputTokens = previousInfo.inputTokens
+    totalOutputTokens = previousInfo.outputTokens
+    totalReasoningTokens = previousInfo.reasoningTokens
+    totalCost = previousInfo.cost
+    actualUsageTokens = previousInfo.actualUsageTokens
+  }
+
+  const instanceProviders = providers().get(instanceId) || []
+
+  const sessionModel = session.model
+  const sessionProviderId = sessionModel?.providerId
+  const sessionModelId = sessionModel?.modelId
+
+  const latestInfo = latestEntry?.messageId ? store.getMessageInfo(latestEntry.messageId) : undefined
+  const latestProviderId = (latestInfo as any)?.providerID || (latestInfo as any)?.providerId || ""
+  const latestModelId = (latestInfo as any)?.modelID || (latestInfo as any)?.modelId || ""
+
+  const selectedModel =
+    resolveSelectedModel(instanceProviders, sessionProviderId, sessionModelId) ??
+    resolveSelectedModel(instanceProviders, latestProviderId, latestModelId)
+
+  let modelOutputLimit = DEFAULT_MODEL_OUTPUT_LIMIT
+
+  if (selectedModel) {
+    contextWindow = selectedModel.limit?.context ?? 0
+    const outputLimit = selectedModel.limit?.output
+    if (typeof outputLimit === "number" && outputLimit > 0) {
+      modelOutputLimit = Math.min(outputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+    }
+    if ((selectedModel.cost?.input ?? 0) === 0 && (selectedModel.cost?.output ?? 0) === 0) {
+      isSubscriptionModel = true
+    }
+  }
+
+  if (contextWindow === 0 && previousInfo) {
+    contextWindow = previousInfo.contextWindow
+  }
+
+  modelOutputLimit = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+
+  if (previousInfo) {
+    const previousContextWindow = previousInfo.contextWindow
+    const previousContextAvailable = previousInfo.contextAvailableTokens ?? null
+    const previousHasContextUsage = previousContextAvailable !== null && previousContextWindow > 0
+      ? previousContextAvailable < previousContextWindow
+      : false
+
+    if (contextWindow !== previousContextWindow) {
+      contextAvailableTokens = null
+      contextAvailableFromPrevious = false
+      latestHasContextUsage = previousHasContextUsage
+    } else {
+      contextAvailableTokens = previousContextAvailable
+      contextAvailableFromPrevious = true
+      latestHasContextUsage = previousHasContextUsage
+    }
+
+    if (!hasUsageEntries) {
+      isSubscriptionModel = previousInfo.isSubscriptionModel
+    } else if (!isSubscriptionModel) {
+      isSubscriptionModel = previousInfo.isSubscriptionModel
+    }
+  }
+
+  const outputBudget = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+
+  if (!contextAvailableFromPrevious) {
+    if (contextWindow > 0) {
+      if (latestHasContextUsage && actualUsageTokens > 0) {
+        contextAvailableTokens = Math.max(contextWindow - (actualUsageTokens + outputBudget), 0)
+      } else {
+        contextAvailableTokens = contextWindow
+      }
+    } else {
+      contextAvailableTokens = null
+    }
+  }
+
+  setSessionInfoByInstance((prev) => {
+    const next = new Map(prev)
+    const instanceInfo = new Map(prev.get(instanceId))
+    instanceInfo.set(sessionId, {
+      cost: totalCost,
+      contextWindow,
+      isSubscriptionModel,
+      inputTokens: totalInputTokens,
+      outputTokens: totalOutputTokens,
+      reasoningTokens: totalReasoningTokens,
+      actualUsageTokens,
+      modelOutputLimit,
+      contextAvailableTokens,
+    })
+    next.set(instanceId, instanceInfo)
+    return next
+  })
+}