5067db3dd0
Wire message.part.delta SSE events into the v2 message store and append deltas onto existing part fields.
1169 lines
36 KiB
TypeScript
1169 lines
36 KiB
TypeScript
import { batch } from "solid-js"
|
|
import { createStore, produce, reconcile } from "solid-js/store"
|
|
import type { SetStoreFunction } from "solid-js/store"
|
|
import { getLogger } from "../../lib/logger"
|
|
import type { ClientPart, MessageInfo } from "../../types/message"
|
|
import { clearRecordDisplayCacheForMessages } from "./record-display-cache"
|
|
import type {
|
|
InstanceMessageState,
|
|
LatestTodoSnapshot,
|
|
MessageRecord,
|
|
MessageUpsertInput,
|
|
PartUpdateInput,
|
|
PendingPartEntry,
|
|
PermissionEntry,
|
|
QuestionEntry,
|
|
ReplaceMessageIdOptions,
|
|
ScrollSnapshot,
|
|
SessionRecord,
|
|
SessionUpsertInput,
|
|
SessionUsageState,
|
|
UsageEntry,
|
|
} from "./types"
|
|
|
|
const storeLog = getLogger("session")
|
|
|
|
interface MessageStoreHooks {
|
|
onSessionCleared?: (instanceId: string, sessionId: string) => void
|
|
}
|
|
|
|
function createInitialState(instanceId: string): InstanceMessageState {
|
|
return {
|
|
instanceId,
|
|
sessions: {},
|
|
sessionOrder: [],
|
|
messages: {},
|
|
messageInfoVersion: {},
|
|
pendingParts: {},
|
|
sessionRevisions: {},
|
|
permissions: {
|
|
queue: [],
|
|
active: null,
|
|
byMessage: {},
|
|
},
|
|
questions: {
|
|
queue: [],
|
|
active: null,
|
|
byMessage: {},
|
|
},
|
|
usage: {},
|
|
scrollState: {},
|
|
latestTodos: {},
|
|
}
|
|
}
|
|
|
|
function ensurePartId(messageId: string, part: ClientPart, index: number): string {
|
|
if (typeof part.id === "string" && part.id.length > 0) {
|
|
return part.id
|
|
}
|
|
|
|
if (part.type === "tool") {
|
|
throw new Error("Tool part missing id")
|
|
}
|
|
|
|
const fallbackId = `${messageId}-part-${index}`
|
|
part.id = fallbackId
|
|
return fallbackId
|
|
}
|
|
|
|
const PENDING_PART_MAX_AGE_MS = 30_000
|
|
|
|
function clonePart(part: ClientPart): ClientPart {
|
|
// Cloning is intentionally disabled; message parts
|
|
// are stored as received from the backend.
|
|
return part
|
|
}
|
|
|
|
function cloneStructuredValue<T>(value: T): T {
|
|
// Legacy helper kept as a no-op to avoid deep copies.
|
|
return value
|
|
}
|
|
|
|
function areMessageIdListsEqual(a: string[], b: string[]): boolean {
|
|
if (a.length !== b.length) {
|
|
return false
|
|
}
|
|
for (let index = 0; index < a.length; index++) {
|
|
if (a[index] !== b[index]) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
function createEmptyUsageState(): SessionUsageState {
|
|
return {
|
|
entries: {},
|
|
totalInputTokens: 0,
|
|
totalOutputTokens: 0,
|
|
totalReasoningTokens: 0,
|
|
totalCost: 0,
|
|
actualUsageTokens: 0,
|
|
latestMessageId: undefined,
|
|
}
|
|
}
|
|
|
|
function extractUsageEntry(info: MessageInfo | undefined): UsageEntry | null {
|
|
if (!info || info.role !== "assistant") return null
|
|
const messageId = typeof info.id === "string" ? info.id : undefined
|
|
if (!messageId) return null
|
|
const tokens = info.tokens
|
|
if (!tokens) return null
|
|
const inputTokens = tokens.input ?? 0
|
|
const outputTokens = tokens.output ?? 0
|
|
const reasoningTokens = tokens.reasoning ?? 0
|
|
const cacheReadTokens = tokens.cache?.read ?? 0
|
|
const cacheWriteTokens = tokens.cache?.write ?? 0
|
|
if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0 && cacheReadTokens === 0 && cacheWriteTokens === 0) {
|
|
return null
|
|
}
|
|
const combinedTokens = info.summary ? outputTokens : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
|
|
return {
|
|
messageId,
|
|
inputTokens,
|
|
outputTokens,
|
|
reasoningTokens,
|
|
cacheReadTokens,
|
|
cacheWriteTokens,
|
|
combinedTokens,
|
|
cost: info.cost ?? 0,
|
|
timestamp: info.time?.created ?? 0,
|
|
hasContextUsage: inputTokens + cacheReadTokens + cacheWriteTokens > 0,
|
|
}
|
|
}
|
|
|
|
function applyUsageState(state: SessionUsageState, entry: UsageEntry | null) {
|
|
if (!entry) return
|
|
state.entries[entry.messageId] = entry
|
|
state.totalInputTokens += entry.inputTokens
|
|
state.totalOutputTokens += entry.outputTokens
|
|
state.totalReasoningTokens += entry.reasoningTokens
|
|
state.totalCost += entry.cost
|
|
if (!state.latestMessageId || entry.timestamp >= (state.entries[state.latestMessageId]?.timestamp ?? 0)) {
|
|
state.latestMessageId = entry.messageId
|
|
state.actualUsageTokens = entry.combinedTokens
|
|
}
|
|
}
|
|
|
|
function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) {
|
|
if (!messageId) return
|
|
const existing = state.entries[messageId]
|
|
if (!existing) return
|
|
state.totalInputTokens -= existing.inputTokens
|
|
state.totalOutputTokens -= existing.outputTokens
|
|
state.totalReasoningTokens -= existing.reasoningTokens
|
|
state.totalCost -= existing.cost
|
|
delete state.entries[messageId]
|
|
if (state.latestMessageId === messageId) {
|
|
state.latestMessageId = undefined
|
|
state.actualUsageTokens = 0
|
|
let latest: UsageEntry | null = null
|
|
for (const candidate of Object.values(state.entries) as UsageEntry[]) {
|
|
if (!latest || candidate.timestamp >= latest.timestamp) {
|
|
latest = candidate
|
|
}
|
|
}
|
|
if (latest) {
|
|
state.latestMessageId = latest.messageId
|
|
state.actualUsageTokens = latest.combinedTokens
|
|
}
|
|
}
|
|
}
|
|
|
|
function rebuildUsageStateFromInfos(infos: Iterable<MessageInfo>): SessionUsageState {
|
|
const usageState = createEmptyUsageState()
|
|
for (const info of infos) {
|
|
const entry = extractUsageEntry(info)
|
|
if (entry) {
|
|
applyUsageState(usageState, entry)
|
|
}
|
|
}
|
|
return usageState
|
|
}
|
|
|
|
export interface InstanceMessageStore {
|
|
instanceId: string
|
|
state: InstanceMessageState
|
|
setState: SetStoreFunction<InstanceMessageState>
|
|
addOrUpdateSession: (input: SessionUpsertInput) => void
|
|
hydrateMessages: (sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) => void
|
|
upsertMessage: (input: MessageUpsertInput) => void
|
|
applyPartUpdate: (input: PartUpdateInput) => void
|
|
applyPartDelta: (input: { messageId: string; partId: string; field: string; delta: string; bumpRevision?: boolean }) => void
|
|
removeMessage: (messageId: string) => void
|
|
removeMessagePart: (messageId: string, partId: string) => void
|
|
bufferPendingPart: (entry: PendingPartEntry) => void
|
|
flushPendingParts: (messageId: string) => void
|
|
replaceMessageId: (options: ReplaceMessageIdOptions) => void
|
|
setMessageInfo: (messageId: string, info: MessageInfo) => void
|
|
getMessageInfo: (messageId: string) => MessageInfo | undefined
|
|
upsertPermission: (entry: PermissionEntry) => void
|
|
removePermission: (permissionId: string) => void
|
|
getPermissionState: (messageId?: string, partId?: string) => { entry: PermissionEntry; active: boolean } | null
|
|
upsertQuestion: (entry: QuestionEntry) => void
|
|
removeQuestion: (requestId: string) => void
|
|
getQuestionState: (messageId?: string, partId?: string) => { entry: QuestionEntry; active: boolean } | null
|
|
setSessionRevert: (sessionId: string, revert?: SessionRecord["revert"] | null) => void
|
|
getSessionRevert: (sessionId: string) => SessionRecord["revert"] | undefined | null
|
|
rebuildUsage: (sessionId: string, infos: Iterable<MessageInfo>) => void
|
|
getSessionUsage: (sessionId: string) => SessionUsageState | undefined
|
|
setScrollSnapshot: (sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) => void
|
|
getScrollSnapshot: (sessionId: string, scope: string) => ScrollSnapshot | undefined
|
|
getSessionRevision: (sessionId: string) => number
|
|
getSessionMessageIds: (sessionId: string) => string[]
|
|
getMessage: (messageId: string) => MessageRecord | undefined
|
|
getLatestTodoSnapshot: (sessionId: string) => LatestTodoSnapshot | undefined
|
|
clearSession: (sessionId: string) => void
|
|
clearInstance: () => void
|
|
}
|
|
|
|
export function createInstanceMessageStore(instanceId: string, hooks?: MessageStoreHooks): InstanceMessageStore {
|
|
const [state, setState] = createStore<InstanceMessageState>(createInitialState(instanceId))
|
|
|
|
const TODO_TOOL_NAME = "todowrite"
|
|
|
|
const messageInfoCache = new Map<string, MessageInfo>()
|
|
|
|
function isCompletedTodoPart(part: ClientPart | undefined): boolean {
|
|
if (!part || (part as any).type !== "tool") {
|
|
return false
|
|
}
|
|
const toolName = typeof (part as any).tool === "string" ? (part as any).tool : ""
|
|
if (toolName !== TODO_TOOL_NAME) {
|
|
return false
|
|
}
|
|
const toolState = (part as any).state
|
|
if (!toolState || typeof toolState !== "object") {
|
|
return false
|
|
}
|
|
return (toolState as { status?: string }).status === "completed"
|
|
}
|
|
|
|
function recordLatestTodoSnapshot(sessionId: string, snapshot: LatestTodoSnapshot) {
|
|
if (!sessionId) return
|
|
setState("latestTodos", sessionId, (existing) => {
|
|
if (existing && existing.timestamp > snapshot.timestamp) {
|
|
return existing
|
|
}
|
|
return snapshot
|
|
})
|
|
}
|
|
|
|
function maybeUpdateLatestTodoFromRecord(record: MessageRecord | undefined) {
|
|
if (!record || !Array.isArray(record.partIds) || record.partIds.length === 0) {
|
|
return
|
|
}
|
|
for (let index = record.partIds.length - 1; index >= 0; index -= 1) {
|
|
const partId = record.partIds[index]
|
|
const partRecord = record.parts[partId]
|
|
if (!partRecord) continue
|
|
if (isCompletedTodoPart(partRecord.data)) {
|
|
const timestamp = typeof record.updatedAt === "number" ? record.updatedAt : Date.now()
|
|
recordLatestTodoSnapshot(record.sessionId, { messageId: record.id, partId, timestamp })
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
function clearLatestTodoSnapshot(sessionId: string) {
|
|
setState("latestTodos", sessionId, undefined)
|
|
}
|
|
|
|
function bumpSessionRevision(sessionId: string) {
|
|
if (!sessionId) return
|
|
setState("sessionRevisions", sessionId, (value = 0) => value + 1)
|
|
}
|
|
|
|
function getSessionRevisionValue(sessionId: string) {
|
|
return state.sessionRevisions[sessionId] ?? 0
|
|
}
|
|
|
|
function withUsageState(sessionId: string, updater: (draft: SessionUsageState) => void) {
|
|
setState("usage", sessionId, (current) => {
|
|
const draft = current
|
|
? {
|
|
...current,
|
|
entries: { ...current.entries },
|
|
}
|
|
: createEmptyUsageState()
|
|
updater(draft)
|
|
return draft
|
|
})
|
|
}
|
|
|
|
function updateUsageWithInfo(info: MessageInfo | undefined) {
|
|
if (!info || typeof info.sessionID !== "string") return
|
|
const messageId = typeof info.id === "string" ? info.id : undefined
|
|
if (!messageId) return
|
|
withUsageState(info.sessionID, (draft) => {
|
|
removeUsageEntry(draft, messageId)
|
|
const entry = extractUsageEntry(info)
|
|
if (entry) {
|
|
applyUsageState(draft, entry)
|
|
}
|
|
})
|
|
}
|
|
|
|
function rebuildUsage(sessionId: string, infos: Iterable<MessageInfo>) {
|
|
const usageState = rebuildUsageStateFromInfos(infos)
|
|
setState("usage", sessionId, usageState)
|
|
}
|
|
|
|
function getSessionUsage(sessionId: string) {
|
|
return state.usage[sessionId]
|
|
}
|
|
|
|
function ensureSessionEntry(sessionId: string): SessionRecord {
|
|
const existing = state.sessions[sessionId]
|
|
if (existing) {
|
|
return existing
|
|
}
|
|
|
|
const now = Date.now()
|
|
const session: SessionRecord = {
|
|
id: sessionId,
|
|
createdAt: now,
|
|
updatedAt: now,
|
|
messageIds: [],
|
|
}
|
|
|
|
setState("sessions", sessionId, session)
|
|
setState("sessionOrder", (order) => (order.includes(sessionId) ? order : [...order, sessionId]))
|
|
return session
|
|
}
|
|
|
|
function addOrUpdateSession(input: SessionUpsertInput) {
|
|
const session = ensureSessionEntry(input.id)
|
|
const previousIds = [...session.messageIds]
|
|
const nextMessageIds = Array.isArray(input.messageIds) ? input.messageIds : session.messageIds
|
|
|
|
setState("sessions", input.id, {
|
|
...session,
|
|
title: input.title ?? session.title,
|
|
parentId: input.parentId ?? session.parentId ?? null,
|
|
updatedAt: Date.now(),
|
|
messageIds: nextMessageIds,
|
|
revert: input.revert ?? session.revert ?? null,
|
|
})
|
|
|
|
if (Array.isArray(input.messageIds) && !areMessageIdListsEqual(previousIds, nextMessageIds)) {
|
|
bumpSessionRevision(input.id)
|
|
}
|
|
}
|
|
|
|
function hydrateMessages(sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) {
|
|
if (!Array.isArray(inputs) || inputs.length === 0) return
|
|
|
|
ensureSessionEntry(sessionId)
|
|
|
|
const incomingIds = inputs.map((item) => item.id)
|
|
|
|
const normalizedRecords: Record<string, MessageRecord> = {}
|
|
const now = Date.now()
|
|
|
|
inputs.forEach((input) => {
|
|
const normalizedParts = normalizeParts(input.id, input.parts)
|
|
const shouldBump = Boolean(input.bumpRevision || normalizedParts)
|
|
const previous = state.messages[input.id]
|
|
normalizedRecords[input.id] = {
|
|
id: input.id,
|
|
sessionId: input.sessionId,
|
|
role: input.role,
|
|
status: input.status,
|
|
createdAt: input.createdAt ?? previous?.createdAt ?? now,
|
|
updatedAt: input.updatedAt ?? now,
|
|
isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false,
|
|
revision: previous ? previous.revision + (shouldBump ? 1 : 0) : 0,
|
|
partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
|
|
parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
|
|
}
|
|
})
|
|
|
|
const infoList = infos ? Array.from(infos) : undefined
|
|
const usageState = infoList ? rebuildUsageStateFromInfos(infoList) : state.usage[sessionId]
|
|
|
|
const nextMessages: Record<string, MessageRecord> = { ...state.messages }
|
|
const nextMessageInfoVersion: Record<string, number> = { ...state.messageInfoVersion }
|
|
const nextPendingParts: Record<string, PendingPartEntry[]> = { ...state.pendingParts }
|
|
const nextPermissionsByMessage: Record<string, Record<string, PermissionEntry>> = {
|
|
...state.permissions.byMessage,
|
|
}
|
|
|
|
Object.entries(normalizedRecords).forEach(([id, record]) => {
|
|
nextMessages[id] = record
|
|
})
|
|
|
|
if (infoList) {
|
|
for (const info of infoList) {
|
|
const messageId = info.id as string
|
|
messageInfoCache.set(messageId, info)
|
|
const currentVersion = nextMessageInfoVersion[messageId] ?? 0
|
|
nextMessageInfoVersion[messageId] = currentVersion + 1
|
|
}
|
|
}
|
|
|
|
batch(() => {
|
|
setState("messages", () => nextMessages)
|
|
setState("messageInfoVersion", () => nextMessageInfoVersion)
|
|
setState("pendingParts", () => nextPendingParts)
|
|
setState("permissions", "byMessage", () => nextPermissionsByMessage)
|
|
|
|
if (usageState) {
|
|
setState("usage", sessionId, usageState)
|
|
}
|
|
|
|
setState("sessions", sessionId, (session) => ({
|
|
...session,
|
|
messageIds: incomingIds,
|
|
updatedAt: Date.now(),
|
|
}))
|
|
|
|
Object.values(normalizedRecords).forEach((record) => {
|
|
maybeUpdateLatestTodoFromRecord(record)
|
|
})
|
|
|
|
bumpSessionRevision(sessionId)
|
|
})
|
|
}
|
|
|
|
function insertMessageIntoSession(sessionId: string, messageId: string) {
|
|
ensureSessionEntry(sessionId)
|
|
setState("sessions", sessionId, "messageIds", (ids = []) => {
|
|
if (ids.includes(messageId)) {
|
|
return ids
|
|
}
|
|
return [...ids, messageId]
|
|
})
|
|
}
|
|
|
|
function normalizeParts(messageId: string, parts: ClientPart[] | undefined) {
|
|
if (!parts || parts.length === 0) {
|
|
return null
|
|
}
|
|
const map: MessageRecord["parts"] = {}
|
|
const ids: string[] = []
|
|
|
|
parts.forEach((part, index) => {
|
|
const id = ensurePartId(messageId, part, index)
|
|
const cloned = clonePart(part)
|
|
map[id] = {
|
|
id,
|
|
data: cloned,
|
|
revision: 0,
|
|
}
|
|
ids.push(id)
|
|
})
|
|
|
|
return { map, ids }
|
|
}
|
|
|
|
function upsertMessage(input: MessageUpsertInput) {
|
|
const normalizedParts = normalizeParts(input.id, input.parts)
|
|
const shouldBump = Boolean(input.bumpRevision || normalizedParts)
|
|
const now = Date.now()
|
|
|
|
let nextRecord: MessageRecord | undefined
|
|
|
|
setState("messages", input.id, (previous) => {
|
|
const revision = previous ? previous.revision + (shouldBump ? 1 : 0) : 0
|
|
const record: MessageRecord = {
|
|
id: input.id,
|
|
sessionId: input.sessionId,
|
|
role: input.role,
|
|
status: input.status,
|
|
createdAt: input.createdAt ?? previous?.createdAt ?? now,
|
|
updatedAt: input.updatedAt ?? now,
|
|
isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false,
|
|
revision,
|
|
partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
|
|
parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
|
|
}
|
|
nextRecord = record
|
|
return record
|
|
})
|
|
|
|
if (nextRecord) {
|
|
maybeUpdateLatestTodoFromRecord(nextRecord)
|
|
}
|
|
|
|
insertMessageIntoSession(input.sessionId, input.id)
|
|
flushPendingParts(input.id)
|
|
bumpSessionRevision(input.sessionId)
|
|
}
|
|
|
|
function bufferPendingPart(entry: PendingPartEntry) {
|
|
setState("pendingParts", entry.messageId, (list = []) => [...list, entry])
|
|
}
|
|
|
|
function clearPendingPartsForMessage(messageId: string) {
|
|
setState("pendingParts", (prev) => {
|
|
if (!prev[messageId]) {
|
|
return prev
|
|
}
|
|
const next = { ...prev }
|
|
delete next[messageId]
|
|
return next
|
|
})
|
|
}
|
|
|
|
function rebindPermissionForPart(messageId: string, partId: string, part: ClientPart) {
|
|
if (!messageId || !partId || part.type !== "tool") {
|
|
return
|
|
}
|
|
|
|
const toolCallId =
|
|
(part as any).callID ??
|
|
(part as any).callId ??
|
|
(part as any).toolCallID ??
|
|
(part as any).toolCallId ??
|
|
undefined
|
|
if (!toolCallId) {
|
|
return
|
|
}
|
|
|
|
setState(
|
|
"permissions",
|
|
"byMessage",
|
|
messageId,
|
|
produce((draft) => {
|
|
if (!draft) return
|
|
const existing = draft[partId]
|
|
for (const [key, entry] of Object.entries(draft)) {
|
|
if (!entry || entry.partId) continue
|
|
const permissionCallId =
|
|
(entry.permission as any).tool?.callID ??
|
|
(entry.permission as any).tool?.callId ??
|
|
(entry.permission as any).callID ??
|
|
(entry.permission as any).callId ??
|
|
(entry.permission as any).toolCallID ??
|
|
(entry.permission as any).toolCallId ??
|
|
(entry.permission as any).metadata?.callID ??
|
|
(entry.permission as any).metadata?.callId ??
|
|
undefined
|
|
if (permissionCallId !== toolCallId) continue
|
|
if (!existing || existing.permission.id === entry.permission.id) {
|
|
entry.partId = partId
|
|
draft[partId] = entry
|
|
delete draft[key]
|
|
}
|
|
break
|
|
}
|
|
}),
|
|
)
|
|
}
|
|
|
|
function applyPartUpdate(input: PartUpdateInput) {
|
|
const message = state.messages[input.messageId]
|
|
if (!message) {
|
|
bufferPendingPart({ messageId: input.messageId, part: input.part, receivedAt: Date.now() })
|
|
return
|
|
}
|
|
|
|
const partId = ensurePartId(input.messageId, input.part, message.partIds.length)
|
|
const cloned = clonePart(input.part)
|
|
|
|
setState(
|
|
"messages",
|
|
input.messageId,
|
|
produce((draft: MessageRecord) => {
|
|
if (!draft.partIds.includes(partId)) {
|
|
draft.partIds = [...draft.partIds, partId]
|
|
}
|
|
const existing = draft.parts[partId]
|
|
const nextRevision = existing ? existing.revision + 1 : (cloned as any).version ?? 0
|
|
draft.parts[partId] = {
|
|
id: partId,
|
|
data: cloned,
|
|
revision: nextRevision,
|
|
}
|
|
draft.updatedAt = Date.now()
|
|
if (input.bumpRevision ?? true) {
|
|
draft.revision += 1
|
|
}
|
|
}),
|
|
)
|
|
|
|
rebindPermissionForPart(input.messageId, partId, cloned)
|
|
|
|
if (isCompletedTodoPart(cloned)) {
|
|
recordLatestTodoSnapshot(message.sessionId, {
|
|
messageId: input.messageId,
|
|
partId,
|
|
timestamp: Date.now(),
|
|
})
|
|
}
|
|
|
|
// Any part update can change the rendered height of the message
|
|
// list, so we treat it as a session revision for scroll purposes.
|
|
bumpSessionRevision(message.sessionId)
|
|
}
|
|
|
|
function applyPartDelta(input: { messageId: string; partId: string; field: string; delta: string; bumpRevision?: boolean }) {
|
|
if (!input?.messageId || !input.partId || !input.field || typeof input.delta !== "string") {
|
|
return
|
|
}
|
|
|
|
const message = state.messages[input.messageId]
|
|
if (!message) {
|
|
// Best-effort: drop deltas for unknown messages.
|
|
return
|
|
}
|
|
|
|
let applied = false
|
|
|
|
setState(
|
|
"messages",
|
|
input.messageId,
|
|
produce((draft: MessageRecord) => {
|
|
const entry = draft.parts[input.partId]
|
|
if (!entry?.data) return
|
|
const part = entry.data as any
|
|
const currentValue = part?.[input.field]
|
|
if (typeof currentValue === "string" || currentValue === undefined || currentValue === null) {
|
|
part[input.field] = `${currentValue ?? ""}${input.delta}`
|
|
applied = true
|
|
}
|
|
if (!applied) return
|
|
entry.revision += 1
|
|
draft.updatedAt = Date.now()
|
|
if (input.bumpRevision ?? true) {
|
|
draft.revision += 1
|
|
}
|
|
}),
|
|
)
|
|
|
|
if (applied) {
|
|
bumpSessionRevision(message.sessionId)
|
|
}
|
|
}
|
|
|
|
function removeMessage(messageId: string) {
|
|
if (!messageId) return
|
|
|
|
const record = state.messages[messageId]
|
|
const sessionIds = new Set<string>()
|
|
|
|
if (record?.sessionId) {
|
|
sessionIds.add(record.sessionId)
|
|
} else {
|
|
Object.values(state.sessions).forEach((session) => {
|
|
if (session.messageIds.includes(messageId)) {
|
|
sessionIds.add(session.id)
|
|
}
|
|
})
|
|
}
|
|
|
|
clearRecordDisplayCacheForMessages(instanceId, [messageId])
|
|
|
|
batch(() => {
|
|
sessionIds.forEach((sessionId) => {
|
|
setState("sessions", sessionId, "messageIds", (ids = []) => ids.filter((id) => id !== messageId))
|
|
})
|
|
|
|
setState("messages", (prev) => {
|
|
if (!prev[messageId]) return prev
|
|
const next = { ...prev }
|
|
delete next[messageId]
|
|
return next
|
|
})
|
|
|
|
setState("messageInfoVersion", (prev) => {
|
|
if (!(messageId in prev)) return prev
|
|
const next = { ...prev }
|
|
delete next[messageId]
|
|
return next
|
|
})
|
|
|
|
messageInfoCache.delete(messageId)
|
|
|
|
setState("pendingParts", (prev) => {
|
|
if (!prev[messageId]) return prev
|
|
const next = { ...prev }
|
|
delete next[messageId]
|
|
return next
|
|
})
|
|
|
|
setState("permissions", "byMessage", (prev) => {
|
|
if (!prev[messageId]) return prev
|
|
const next = { ...prev }
|
|
delete next[messageId]
|
|
return next
|
|
})
|
|
|
|
sessionIds.forEach((sessionId) => {
|
|
withUsageState(sessionId, (draft) => removeUsageEntry(draft, messageId))
|
|
if (state.latestTodos[sessionId]?.messageId === messageId) {
|
|
clearLatestTodoSnapshot(sessionId)
|
|
}
|
|
bumpSessionRevision(sessionId)
|
|
})
|
|
})
|
|
}
|
|
|
|
function removeMessagePart(messageId: string, partId: string) {
|
|
if (!messageId || !partId) return
|
|
const message = state.messages[messageId]
|
|
if (!message) return
|
|
|
|
clearRecordDisplayCacheForMessages(instanceId, [messageId])
|
|
|
|
batch(() => {
|
|
setState(
|
|
"messages",
|
|
messageId,
|
|
produce((draft: MessageRecord) => {
|
|
if (!draft.parts[partId] && !draft.partIds.includes(partId)) return
|
|
draft.partIds = draft.partIds.filter((id) => id !== partId)
|
|
delete draft.parts[partId]
|
|
draft.updatedAt = Date.now()
|
|
draft.revision += 1
|
|
}),
|
|
)
|
|
|
|
setState("permissions", "byMessage", messageId, (prev) => {
|
|
if (!prev || !prev[partId]) return prev
|
|
const next = { ...prev }
|
|
delete next[partId]
|
|
return next
|
|
})
|
|
|
|
bumpSessionRevision(message.sessionId)
|
|
})
|
|
}
|
|
|
|
|
|
function flushPendingParts(messageId: string) {
|
|
const pending = state.pendingParts[messageId]
|
|
if (!pending || pending.length === 0) {
|
|
return
|
|
}
|
|
const now = Date.now()
|
|
const validEntries = pending.filter((entry) => now - entry.receivedAt <= PENDING_PART_MAX_AGE_MS)
|
|
if (validEntries.length === 0) {
|
|
clearPendingPartsForMessage(messageId)
|
|
return
|
|
}
|
|
validEntries.forEach((entry) => applyPartUpdate({ messageId, part: entry.part }))
|
|
clearPendingPartsForMessage(messageId)
|
|
}
|
|
|
|
function replaceMessageId(options: ReplaceMessageIdOptions) {
|
|
if (options.oldId === options.newId) return
|
|
const existing = state.messages[options.oldId]
|
|
if (!existing) return
|
|
|
|
const cloned: MessageRecord = {
|
|
...existing,
|
|
id: options.newId,
|
|
isEphemeral: false,
|
|
updatedAt: Date.now(),
|
|
}
|
|
|
|
setState("messages", options.newId, cloned)
|
|
setState("messages", (prev) => {
|
|
const next = { ...prev }
|
|
delete next[options.oldId]
|
|
return next
|
|
})
|
|
|
|
const affectedSessions = new Set<string>()
|
|
|
|
Object.values(state.sessions).forEach((session) => {
|
|
const index = session.messageIds.indexOf(options.oldId)
|
|
if (index === -1) return
|
|
setState("sessions", session.id, "messageIds", (ids) => {
|
|
const next = [...ids]
|
|
next[index] = options.newId
|
|
return next
|
|
})
|
|
affectedSessions.add(session.id)
|
|
})
|
|
|
|
affectedSessions.forEach((sessionId) => bumpSessionRevision(sessionId))
|
|
|
|
const infoEntry = messageInfoCache.get(options.oldId)
|
|
if (infoEntry) {
|
|
messageInfoCache.set(options.newId, infoEntry)
|
|
messageInfoCache.delete(options.oldId)
|
|
const version = state.messageInfoVersion[options.oldId] ?? 0
|
|
setState("messageInfoVersion", options.newId, version)
|
|
setState("messageInfoVersion", (prev) => {
|
|
const next = { ...prev }
|
|
delete next[options.oldId]
|
|
return next
|
|
})
|
|
}
|
|
|
|
const permissionMap = state.permissions.byMessage[options.oldId]
|
|
if (permissionMap) {
|
|
setState("permissions", "byMessage", options.newId, permissionMap)
|
|
setState("permissions", (prev) => {
|
|
const next = { ...prev }
|
|
const nextByMessage = { ...next.byMessage }
|
|
delete nextByMessage[options.oldId]
|
|
next.byMessage = nextByMessage
|
|
return next
|
|
})
|
|
}
|
|
|
|
const questionMap = state.questions.byMessage[options.oldId]
|
|
if (questionMap) {
|
|
setState("questions", "byMessage", options.newId, questionMap)
|
|
setState("questions", (prev) => {
|
|
const next = { ...prev }
|
|
const nextByMessage = { ...next.byMessage }
|
|
delete nextByMessage[options.oldId]
|
|
next.byMessage = nextByMessage
|
|
return next
|
|
})
|
|
}
|
|
|
|
const pending = state.pendingParts[options.oldId]
|
|
if (pending) {
|
|
setState("pendingParts", options.newId, pending)
|
|
}
|
|
clearPendingPartsForMessage(options.oldId)
|
|
maybeUpdateLatestTodoFromRecord(cloned)
|
|
}
|
|
|
|
function setMessageInfo(messageId: string, info: MessageInfo) {
|
|
if (!messageId) return
|
|
messageInfoCache.set(messageId, info)
|
|
const nextVersion = (state.messageInfoVersion[messageId] ?? 0) + 1
|
|
setState("messageInfoVersion", messageId, nextVersion)
|
|
updateUsageWithInfo(info)
|
|
}
|
|
|
|
function getMessageInfo(messageId: string) {
|
|
void state.messageInfoVersion[messageId]
|
|
return messageInfoCache.get(messageId)
|
|
}
|
|
|
|
function upsertPermission(entry: PermissionEntry) {
|
|
const messageKey = entry.messageId ?? "__global__"
|
|
const partKey = entry.partId ?? entry.permission?.id ?? "__global__"
|
|
|
|
setState(
|
|
"permissions",
|
|
produce((draft) => {
|
|
draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {}
|
|
draft.byMessage[messageKey][partKey] = entry
|
|
const existingIndex = draft.queue.findIndex((item) => item.permission.id === entry.permission.id)
|
|
if (existingIndex === -1) {
|
|
draft.queue.push(entry)
|
|
} else {
|
|
draft.queue[existingIndex] = entry
|
|
}
|
|
if (!draft.active || draft.active.permission.id === entry.permission.id) {
|
|
draft.active = entry
|
|
}
|
|
}),
|
|
)
|
|
}
|
|
|
|
function removePermission(permissionId: string) {
|
|
setState(
|
|
"permissions",
|
|
produce((draft) => {
|
|
draft.queue = draft.queue.filter((item) => item.permission.id !== permissionId)
|
|
if (draft.active?.permission.id === permissionId) {
|
|
draft.active = draft.queue[0] ?? null
|
|
}
|
|
Object.keys(draft.byMessage).forEach((messageKey) => {
|
|
const partEntries = draft.byMessage[messageKey]
|
|
Object.keys(partEntries).forEach((partKey) => {
|
|
if (partEntries[partKey].permission.id === permissionId) {
|
|
delete partEntries[partKey]
|
|
}
|
|
})
|
|
if (Object.keys(partEntries).length === 0) {
|
|
delete draft.byMessage[messageKey]
|
|
}
|
|
})
|
|
}),
|
|
)
|
|
}
|
|
|
|
function getPermissionState(messageId?: string, partId?: string) {
|
|
const messageKey = messageId ?? "__global__"
|
|
const partKey = partId ?? "__global__"
|
|
const entry = state.permissions.byMessage[messageKey]?.[partKey]
|
|
if (!entry) return null
|
|
const active = state.permissions.active?.permission.id === entry.permission.id
|
|
return { entry, active }
|
|
}
|
|
|
|
function upsertQuestion(entry: QuestionEntry) {
|
|
const messageKey = entry.messageId ?? "__global__"
|
|
const partKey = entry.partId ?? entry.request?.id ?? "__global__"
|
|
|
|
setState(
|
|
"questions",
|
|
produce((draft) => {
|
|
draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {}
|
|
draft.byMessage[messageKey][partKey] = entry
|
|
const existingIndex = draft.queue.findIndex((item) => item.request.id === entry.request.id)
|
|
if (existingIndex === -1) {
|
|
draft.queue.push(entry)
|
|
} else {
|
|
draft.queue[existingIndex] = entry
|
|
}
|
|
if (!draft.active || draft.active.request.id === entry.request.id) {
|
|
draft.active = entry
|
|
}
|
|
}),
|
|
)
|
|
}
|
|
|
|
function removeQuestion(requestId: string) {
|
|
setState(
|
|
"questions",
|
|
produce((draft) => {
|
|
draft.queue = draft.queue.filter((item) => item.request.id !== requestId)
|
|
if (draft.active?.request.id === requestId) {
|
|
draft.active = draft.queue[0] ?? null
|
|
}
|
|
Object.keys(draft.byMessage).forEach((messageKey) => {
|
|
const partEntries = draft.byMessage[messageKey]
|
|
Object.keys(partEntries).forEach((partKey) => {
|
|
if (partEntries[partKey].request.id === requestId) {
|
|
delete partEntries[partKey]
|
|
}
|
|
})
|
|
if (Object.keys(partEntries).length === 0) {
|
|
delete draft.byMessage[messageKey]
|
|
}
|
|
})
|
|
}),
|
|
)
|
|
}
|
|
|
|
function getQuestionState(messageId?: string, partId?: string) {
|
|
const messageKey = messageId ?? "__global__"
|
|
const partKey = partId ?? "__global__"
|
|
const entry = state.questions.byMessage[messageKey]?.[partKey]
|
|
if (!entry) return null
|
|
const active = state.questions.active?.request.id === entry.request.id
|
|
return { entry, active }
|
|
}
|
|
|
|
function pruneMessagesAfterRevert(sessionId: string, revertMessageId: string) {
|
|
const session = state.sessions[sessionId]
|
|
if (!session) return
|
|
const stopIndex = session.messageIds.indexOf(revertMessageId)
|
|
if (stopIndex === -1) return
|
|
const removedIds = session.messageIds.slice(stopIndex)
|
|
const keptIds = session.messageIds.slice(0, stopIndex)
|
|
if (removedIds.length === 0) return
|
|
|
|
setState("sessions", sessionId, "messageIds", keptIds)
|
|
|
|
setState("messages", (prev) => {
|
|
const next = { ...prev }
|
|
removedIds.forEach((id) => delete next[id])
|
|
return next
|
|
})
|
|
|
|
setState("messageInfoVersion", (prev) => {
|
|
const next = { ...prev }
|
|
removedIds.forEach((id) => delete next[id])
|
|
return next
|
|
})
|
|
|
|
removedIds.forEach((id) => messageInfoCache.delete(id))
|
|
|
|
setState("pendingParts", (prev) => {
|
|
const next = { ...prev }
|
|
removedIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("permissions", "byMessage", (prev) => {
|
|
const next = { ...prev }
|
|
removedIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("questions", "byMessage", (prev) => {
|
|
const next = { ...prev }
|
|
removedIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
withUsageState(sessionId, (draft) => {
|
|
removedIds.forEach((id) => removeUsageEntry(draft, id))
|
|
})
|
|
|
|
bumpSessionRevision(sessionId)
|
|
}
|
|
|
|
function setSessionRevert(sessionId: string, revert?: SessionRecord["revert"] | null) {
|
|
if (!sessionId) return
|
|
ensureSessionEntry(sessionId)
|
|
if (revert?.messageID) {
|
|
pruneMessagesAfterRevert(sessionId, revert.messageID)
|
|
}
|
|
setState("sessions", sessionId, "revert", revert ?? null)
|
|
}
|
|
|
|
function getSessionRevert(sessionId: string) {
|
|
return state.sessions[sessionId]?.revert ?? null
|
|
}
|
|
|
|
function makeScrollKey(sessionId: string, scope: string) {
|
|
return `${sessionId}:${scope}`
|
|
}
|
|
|
|
function setScrollSnapshot(sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) {
|
|
const key = makeScrollKey(sessionId, scope)
|
|
setState("scrollState", key, { ...snapshot, updatedAt: Date.now() })
|
|
}
|
|
|
|
function getScrollSnapshot(sessionId: string, scope: string) {
|
|
const key = makeScrollKey(sessionId, scope)
|
|
return state.scrollState[key]
|
|
}
|
|
|
|
function clearSession(sessionId: string) {
|
|
if (!sessionId) return
|
|
|
|
const messageIds = Object.values(state.messages)
|
|
.filter((record) => record.sessionId === sessionId)
|
|
.map((record) => record.id)
|
|
|
|
storeLog.info("Clearing session data", { instanceId, sessionId, messageCount: messageIds.length })
|
|
clearRecordDisplayCacheForMessages(instanceId, messageIds)
|
|
|
|
batch(() => {
|
|
setState("messages", (prev) => {
|
|
const next = { ...prev }
|
|
messageIds.forEach((id) => delete next[id])
|
|
return next
|
|
})
|
|
|
|
setState("messageInfoVersion", (prev) => {
|
|
const next = { ...prev }
|
|
messageIds.forEach((id) => delete next[id])
|
|
return next
|
|
})
|
|
|
|
messageIds.forEach((id) => messageInfoCache.delete(id))
|
|
|
|
setState("pendingParts", (prev) => {
|
|
const next = { ...prev }
|
|
messageIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("permissions", "byMessage", (prev) => {
|
|
const next = { ...prev }
|
|
messageIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("questions", "byMessage", (prev) => {
|
|
const next = { ...prev }
|
|
messageIds.forEach((id) => {
|
|
if (next[id]) delete next[id]
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("usage", (prev) => {
|
|
const next = { ...prev }
|
|
delete next[sessionId]
|
|
return next
|
|
})
|
|
|
|
setState("sessionRevisions", (prev) => {
|
|
const next = { ...prev }
|
|
delete next[sessionId]
|
|
return next
|
|
})
|
|
|
|
setState("scrollState", (prev) => {
|
|
const next = { ...prev }
|
|
const prefix = `${sessionId}:`
|
|
Object.keys(next).forEach((key) => {
|
|
if (key.startsWith(prefix)) {
|
|
delete next[key]
|
|
}
|
|
})
|
|
return next
|
|
})
|
|
|
|
setState("sessions", sessionId, (current) => {
|
|
if (!current) return current
|
|
return { ...current, messageIds: [] }
|
|
})
|
|
|
|
setState("sessions", (prev) => {
|
|
const next = { ...prev }
|
|
delete next[sessionId]
|
|
return next
|
|
})
|
|
|
|
setState("sessionOrder", (ids) => ids.filter((id) => id !== sessionId))
|
|
})
|
|
|
|
clearLatestTodoSnapshot(sessionId)
|
|
|
|
hooks?.onSessionCleared?.(instanceId, sessionId)
|
|
}
|
|
|
|
|
|
function clearInstance() {
|
|
messageInfoCache.clear()
|
|
setState(reconcile(createInitialState(instanceId)))
|
|
}
|
|
|
|
return {
|
|
|
|
instanceId,
|
|
state,
|
|
setState,
|
|
addOrUpdateSession,
|
|
hydrateMessages,
|
|
upsertMessage,
|
|
applyPartUpdate,
|
|
applyPartDelta,
|
|
removeMessage,
|
|
removeMessagePart,
|
|
bufferPendingPart,
|
|
flushPendingParts,
|
|
replaceMessageId,
|
|
setMessageInfo,
|
|
getMessageInfo,
|
|
upsertPermission,
|
|
removePermission,
|
|
getPermissionState,
|
|
upsertQuestion,
|
|
removeQuestion,
|
|
getQuestionState,
|
|
|
|
setSessionRevert,
|
|
getSessionRevert,
|
|
rebuildUsage,
|
|
getSessionUsage,
|
|
setScrollSnapshot,
|
|
getScrollSnapshot,
|
|
getSessionRevision: getSessionRevisionValue,
|
|
getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
|
|
getMessage: (messageId: string) => state.messages[messageId],
|
|
getLatestTodoSnapshot: (sessionId: string) => state.latestTodos[sessionId],
|
|
clearSession,
|
|
clearInstance,
|
|
}
|
|
}
|
|
|