/* eslint-disable @typescript-eslint/no-explicit-any */
import { create } from "zustand";
import { devtools, persist, createJSONStorage } from "zustand/middleware";
import type { Message, Part } from "@opencode-ai/sdk";
import { opencodeClient } from "@/lib/opencode/client";
import { isExecutionForkMetaText } from "@/lib/messages/executionMeta";
import type { SessionMemoryState, MessageStreamLifecycle, AttachedFile } from "./types/sessionTypes";
import { MEMORY_LIMITS } from "./types/sessionTypes";
import {
    touchStreamingLifecycle,
    removeLifecycleEntries,
    clearLifecycleTimersForIds,
    clearLifecycleCompletionTimer
} from "./utils/streamingUtils";
import { extractTextFromPart, normalizeStreamingPart } from "./utils/messageUtils";
import { getSafeStorage } from "./utils/safeStorage";
import { useFileStore } from "./fileStore";
import { useSessionStore } from "./sessionStore";
import { useContextStore } from "./contextStore";

// Helper function to clean up pending user message metadata
const cleanupPendingUserMessageMeta = (
    currentPending: Map<string, { mode?: string; providerID?: string; modelID?: string }>,
    sessionId: string
): Map<string, { mode?: string; providerID?: string; modelID?: string }> => {
    const nextPending = new Map(currentPending);
    nextPending.delete(sessionId);
    return nextPending;
};

interface QueuedPart {
    sessionId: string;
    messageId: string;
    part: Part;
    role?: string;
    currentSessionId?: string;
}

let batchQueue: QueuedPart[] = [];
let flushTimer: ReturnType<typeof setTimeout> | null = null;

const USER_BATCH_WINDOW_MS = 50;
const COMPACTION_WINDOW_MS = 30_000;

const timeoutRegistry = new Map<string, ReturnType<typeof setTimeout>>();
const lastContentRegistry = new Map<string, string>();
const streamingCooldownTimers = new Map<string, ReturnType<typeof setTimeout>>();

const MIN_SORTABLE_LENGTH = 10;

const extractSortableId = (id: unknown): string | null => {
    if (typeof id !== "string") {
        return null;
    }
    const trimmed = id.trim();
    if (!trimmed) {
        return null;
    }
    const underscoreIndex = trimmed.indexOf("_");
    const candidate = underscoreIndex >= 0 ? trimmed.slice(underscoreIndex + 1) : trimmed;
    if (!candidate || candidate.length < MIN_SORTABLE_LENGTH) {
        return null;
    }
    return candidate;
};

const isIdNewer = (id: string, referenceId: string): boolean => {
    const currentSortable = extractSortableId(id);
    const referenceSortable = extractSortableId(referenceId);
    if (!currentSortable || !referenceSortable) {
        return true;
    }
    if (currentSortable.length !== referenceSortable.length) {
        return true;
    }
    return currentSortable > referenceSortable;
};

const streamDebugEnabled = (): boolean => {
    if (typeof window === "undefined") return false;
    try {
        return window.localStorage.getItem("openchamber_stream_debug") === "1";
    } catch {
        return false;
    }
};

const computePartsTextLength = (parts: Part[] | undefined): number => {
    if (!Array.isArray(parts) || parts.length === 0) {
        return 0;
    }
    return parts.reduce((sum, part) => {
        if (!part || part.type !== "text") {
            return sum;
        }
        const content = (part as Record<string, unknown>).text ?? (part as Record<string, unknown>).content;
        if (typeof content === "string") {
            return sum + content.length;
        }
        return sum;
    }, 0);
};

const hasStopReasonStop = (parts: Part[] | undefined): boolean => {
    if (!Array.isArray(parts)) {
        return false;
    }
    return parts.some(
        (part) => part?.type === "step-finish" && (part as Record<string, unknown>)?.reason === "stop"
    );
};

const getPartKey = (part: Part | undefined): string | undefined => {
    if (!part) {
        return undefined;
    }
    if (typeof part.id === "string" && part.id.length > 0) {
        return part.id;
    }
    if (part.type) {
        const reason = (part as Record<string, unknown>).reason;
        const callId = (part as Record<string, unknown>).callID;
        return `${part.type}-${reason ?? ""}-${callId ?? ""}`;
    }
    return undefined;
};

const ignoredAssistantMessageIds = new Set<string>();

const mergePreferExistingParts = (existing: Part[] = [], incoming: Part[] = []): Part[] => {
    if (!incoming.length) {
        return [...existing];
    }
    const merged = [...existing];
    const existingKeys = new Set(existing.map(getPartKey).filter((key): key is string => Boolean(key)));

    incoming.forEach((part) => {
        if (!part) {
            return;
        }
        const key = getPartKey(part);
        if (key && existingKeys.has(key)) {
            return;
        }
        merged.push(part);
        if (key) {
            existingKeys.add(key);
        }
    });

    return merged;
};

const mergeDuplicateMessage = (
    existing: { info: any; parts: Part[] },
    incoming: { info: any; parts: Part[] }
): { info: any; parts: Part[] } => {
    const existingParts = Array.isArray(existing.parts) ? existing.parts : [];
    const incomingParts = Array.isArray(incoming.parts) ? incoming.parts : [];
    const existingLen = computePartsTextLength(existingParts);
    const incomingLen = computePartsTextLength(incomingParts);
    const existingStop = hasStopReasonStop(existingParts);
    const incomingStop = hasStopReasonStop(incomingParts);

    let parts = incomingParts;
    if (existingStop && existingLen >= incomingLen) {
        parts = mergePreferExistingParts(existingParts, incomingParts);
    } else if (incomingStop && incomingLen >= existingLen) {
        parts = mergePreferExistingParts(incomingParts, existingParts);
    } else if (existingLen >= incomingLen) {
        parts = existingParts;
    }

    return {
        ...incoming,
        info: {
            ...existing.info,
            ...incoming.info,
        },
        parts,
    };
};

const dedupeMessagesById = (messages: { info: any; parts: Part[] }[]) => {
    const deduped: { info: any; parts: Part[] }[] = [];
    const indexById = new Map<string, number>();

    for (const message of messages) {
        const messageId = typeof message?.info?.id === "string" ? message.info.id : null;
        if (!messageId) {
            deduped.push(message);
            continue;
        }
        const existingIndex = indexById.get(messageId);
        if (existingIndex === undefined) {
            indexById.set(messageId, deduped.length);
            deduped.push(message);
            continue;
        }
        deduped[existingIndex] = mergeDuplicateMessage(deduped[existingIndex], message);
    }

    return deduped;
};

const computeMaxTrimmedHeadId = (removed: Array<{ info: any }>, previous?: string): string | undefined => {
    let maxId = previous;
    let maxSortable = previous ? extractSortableId(previous) : null;

    for (const entry of removed) {
        const candidateId = entry?.info?.id;
        const candidateSortable = extractSortableId(candidateId);
        if (!candidateId || !candidateSortable) {
            continue;
        }
        if (!maxSortable || candidateSortable > maxSortable) {
            maxSortable = candidateSortable;
            maxId = candidateId;
        }
    }

    return maxId;
};

const setStreamingIdForSession = (source: Map<string, string | null>, sessionId: string, messageId: string | null) => {
    const existing = source.get(sessionId);
    if (existing === messageId) {
        return source;
    }
    const next = new Map(source);
    if (messageId) {
        next.set(sessionId, messageId);
    } else {
        next.delete(sessionId);
    }
    return next;
};

const upsertMessageSessionIndex = (source: Map<string, string>, messageId: string, sessionId: string) => {
    const existing = source.get(messageId);
    if (existing === sessionId) {
        return source;
    }
    const next = new Map(source);
    next.set(messageId, sessionId);
    return next;
};

const removeMessageSessionIndexEntries = (source: Map<string, string>, ids: Iterable<string>) => {
    const next = new Map(source);
    let mutated = false;
    for (const id of ids) {
        if (next.delete(id)) {
            mutated = true;
        }
    }
    return mutated ? next : source;
};

const collectActiveMessageIdsForSession = (state: MessageState, sessionId: string): Set<string> => {
    const ids = new Set<string>();
    const latest = state.streamingMessageIds.get(sessionId);
    if (latest) {
        ids.add(latest);
    }
    state.messageStreamStates.forEach((_lifecycle, messageId) => {
        if (state.messageSessionIndex.get(messageId) === sessionId) {
            ids.add(messageId);
        }
    });
    return ids;
};

const isMessageStreamingInSession = (state: MessageState, sessionId: string, messageId: string) => {
    if (state.streamingMessageIds.get(sessionId) === messageId) {
        return true;
    }
    return state.messageSessionIndex.get(messageId) === sessionId && state.messageStreamStates.has(messageId);
};

const resolveSessionDirectory = async (sessionId: string | null | undefined): Promise<string | undefined> => {
    if (!sessionId) {
        return undefined;
    }

    try {
        const sessionStore = useSessionStore.getState();
        const metadata = sessionStore.getWorktreeMetadata(sessionId);
        if (metadata?.path) {
            return metadata.path;
        }

        const session = sessionStore.sessions.find((entry) => entry.id === sessionId) as { directory?: string } | undefined;
        const sessionDirectory =
            typeof session?.directory === 'string' && session.directory.length > 0 ? session.directory : undefined;

        return sessionDirectory;
    } catch (error) {
        console.warn('Failed to resolve session directory override:', error);
        return undefined;
    }
};

const getSessionRevertMessageId = (sessionId: string | null | undefined): string | null => {
    if (!sessionId) return null;
    try {
        const sessionStore = useSessionStore.getState();
        const session = sessionStore.sessions.find((entry) => entry.id === sessionId) as { revert?: { messageID?: string } } | undefined;
        return session?.revert?.messageID ?? null;
    } catch {
        return null;
    }
};

const filterRevertedMessages = (
    messages: { info: Message; parts: Part[] }[],
    revertMessageId: string | null
): { info: Message; parts: Part[] }[] => {
    if (!revertMessageId) return messages;

    const revertIndex = messages.findIndex((m) => m.info.id === revertMessageId);
    if (revertIndex === -1) return messages;

    // Keep only messages before the revert point (exclusive)
    return messages.slice(0, revertIndex);
};

const executeWithSessionDirectory = async <T>(sessionId: string | null | undefined, operation: () => Promise<T>): Promise<T> => {
    const directoryOverride = await resolveSessionDirectory(sessionId);
    if (directoryOverride) {
        return opencodeClient.withDirectory(directoryOverride, operation);
    }
    return operation();
};

interface SessionAbortRecord {
    timestamp: number;
    acknowledged: boolean;
}

interface MessageState {
    messages: Map<string, { info: any; parts: Part[] }[]>;
    sessionMemoryState: Map<string, SessionMemoryState>;
    messageStreamStates: Map<string, MessageStreamLifecycle>;
    messageSessionIndex: Map<string, string>;
    streamingMessageIds: Map<string, string | null>;
    abortControllers: Map<string, AbortController>;
    lastUsedProvider: { providerID: string; modelID: string } | null;
    isSyncing: boolean;
    pendingAssistantParts: Map<string, { sessionId: string; parts: Part[] }>;
    sessionCompactionUntil: Map<string, number>;
    sessionAbortFlags: Map<string, SessionAbortRecord>;
    pendingAssistantHeaderSessions: Set<string>;
    pendingUserMessageMetaBySession: Map<string, { mode?: string; providerID?: string; modelID?: string }>;
}

interface MessageActions {
    loadMessages: (sessionId: string) => Promise<void>;
    sendMessage: (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null) => 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;
    forceCompleteMessage: (sessionId: string | null | undefined, messageId: string, source?: "timeout" | "cooldown") => void;
    completeStreamingMessage: (sessionId: string, messageId: string) => void;
    markMessageStreamSettled: (messageId: string) => void;
    updateMessageInfo: (sessionId: string, messageId: string, messageInfo: any) => void;
    syncMessages: (sessionId: string, messages: { info: Message; parts: Part[] }[]) => void;
    updateViewportAnchor: (sessionId: string, anchor: number) => void;
    trimToViewportWindow: (sessionId: string, targetSize?: number, currentSessionId?: string) => void;
    evictLeastRecentlyUsed: (currentSessionId?: string) => void;
    loadMoreMessages: (sessionId: string, direction: "up" | "down") => Promise<void>;
    getLastMessageModel: (sessionId: string) => { providerID?: string; modelID?: string } | null;
    updateSessionCompaction: (sessionId: string, compactingTimestamp: number | null | undefined) => void;
    acknowledgeSessionAbort: (sessionId: string) => void;
}

type MessageStore = MessageState & MessageActions;

export const useMessageStore = create<MessageStore>()(
    devtools(
        persist(
            (set, get) => ({

                messages: new Map(),
                sessionMemoryState: new Map(),
                messageStreamStates: new Map(),
                messageSessionIndex: new Map(),
                streamingMessageIds: new Map(),
                abortControllers: new Map(),
                lastUsedProvider: null,
                isSyncing: false,
                pendingAssistantParts: new Map(),
                sessionCompactionUntil: new Map(),
                sessionAbortFlags: new Map(),
                pendingAssistantHeaderSessions: new Set(),
                pendingUserMessageMetaBySession: new Map(),

                loadMessages: async (sessionId: string, limit: number = MEMORY_LIMITS.VIEWPORT_MESSAGES) => {
                        const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId));

                        // Filter out reverted messages first
                        const revertMessageId = getSessionRevertMessageId(sessionId);
                        const messagesWithoutReverted = filterRevertedMessages(allMessages, revertMessageId);

                        const watermark = get().sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;

                        const afterWatermark = watermark
                            ? messagesWithoutReverted.filter((message) => {
                                  const messageId = message?.info?.id;
                                  if (!messageId) return true;
                                  return isIdNewer(messageId, watermark);
                              })
                            : messagesWithoutReverted;
                        const messagesToKeep = afterWatermark.slice(-limit);

                        set((state) => {
                            const newMessages = new Map(state.messages);
                            const previousMessages = state.messages.get(sessionId) || [];
                            const previousMessagesById = new Map(
                                previousMessages
                                    .filter((msg) => typeof msg.info?.id === "string")
                                    .map((msg) => [msg.info.id as string, msg])
                            );

                            const normalizedMessages = messagesToKeep.map((message) => {
                                const infoWithMarker = {
                                    ...message.info,
                                    clientRole: (message.info as any)?.clientRole ?? message.info.role,
                                    userMessageMarker: message.info.role === "user" ? true : (message.info as any)?.userMessageMarker,
                                } as any;

                                const serverParts = (Array.isArray(message.parts) ? message.parts : []).map((part) => {
                                    if (part?.type === 'text') {
                                        const raw = (part as any).text ?? (part as any).content ?? '';
                                        if (isExecutionForkMetaText(raw)) {
                                            return { ...part, synthetic: true } as Part;
                                        }
                                    }
                                    return part;
                                });
                                const existingEntry = infoWithMarker?.id
                                    ? previousMessagesById.get(infoWithMarker.id as string)
                                    : undefined;

                                if (
                                    existingEntry &&
                                    existingEntry.info.role === "assistant"
                                ) {
                                    const existingParts = Array.isArray(existingEntry.parts) ? existingEntry.parts : [];
                                    const existingLen = computePartsTextLength(existingParts);
                                    const serverLen = computePartsTextLength(serverParts);
                                    const storeHasStop = hasStopReasonStop(existingParts);

                                    if (storeHasStop && existingLen > serverLen) {
                                        const mergedParts = mergePreferExistingParts(existingParts, serverParts);
                                        return {
                                            ...message,
                                            info: infoWithMarker,
                                            parts: mergedParts,
                                        };
                                    }
                                }

                                return {
                                    ...message,
                                    info: infoWithMarker,
                                    parts: serverParts,
                                };
                            });

                            const mergedMessages = dedupeMessagesById(normalizedMessages);

                            const previousIds = new Set(previousMessages.map((msg) => msg.info.id));
                            const nextIds = new Set(mergedMessages.map((msg) => msg.info.id));
                            const removedIds: string[] = [];
                            previousIds.forEach((id) => {
                                if (!nextIds.has(id)) {
                                    removedIds.push(id);
                                }
                            });

                            newMessages.set(sessionId, mergedMessages);

                            const newMemoryState = new Map(state.sessionMemoryState);
                            const previousMemoryState = state.sessionMemoryState.get(sessionId);
                            newMemoryState.set(sessionId, {
                                viewportAnchor: mergedMessages.length - 1,
                                isStreaming: false,
                                lastAccessedAt: Date.now(),
                                backgroundMessageCount: 0,
                                totalAvailableMessages: allMessages.length,
                                hasMoreAbove: allMessages.length > messagesToKeep.length,
                                trimmedHeadMaxId: previousMemoryState?.trimmedHeadMaxId,
                                streamingCooldownUntil: undefined,
                            });

                            const result: Record<string, any> = {
                                messages: newMessages,
                                sessionMemoryState: newMemoryState,
                            };

                        clearLifecycleTimersForIds(removedIds);
                        const updatedLifecycle = removeLifecycleEntries(state.messageStreamStates, removedIds);
                        if (updatedLifecycle !== state.messageStreamStates) {
                            result.messageStreamStates = updatedLifecycle;
                        }

                        if (removedIds.length > 0) {
                            const currentStreaming = state.streamingMessageIds.get(sessionId);
                            if (currentStreaming && removedIds.includes(currentStreaming)) {
                                result.streamingMessageIds = setStreamingIdForSession(
                                    result.streamingMessageIds ?? state.streamingMessageIds,
                                    sessionId,
                                    null
                                );
                            }
                        }

                        if (removedIds.length > 0) {
                            const nextIndex = removeMessageSessionIndexEntries(
                                result.messageSessionIndex ?? state.messageSessionIndex,
                                removedIds
                            );
                            if (nextIndex !== (result.messageSessionIndex ?? state.messageSessionIndex)) {
                                result.messageSessionIndex = nextIndex;
                            }
                        }

                        if (removedIds.length > 0) {
                            const nextPendingParts = new Map(state.pendingAssistantParts);
                            let pendingChanged = false;
                            removedIds.forEach((id) => {
                                if (nextPendingParts.delete(id)) {
                                    pendingChanged = true;
                                }
                            });
                            if (pendingChanged) {
                                result.pendingAssistantParts = nextPendingParts;
                            }
                        }

                        const targetIndex = result.messageSessionIndex ?? state.messageSessionIndex;
                        let indexAccumulator = targetIndex;
                        mergedMessages.forEach((message) => {
                            const id = (message?.info as { id?: unknown })?.id;
                            if (typeof id === "string" && id.length > 0) {
                                indexAccumulator = upsertMessageSessionIndex(indexAccumulator, id, sessionId);
                            }
                        });
                        if (indexAccumulator !== targetIndex) {
                            result.messageSessionIndex = indexAccumulator;
                        }

                        return result;

                        });
                },

                sendMessage: async (content: string, providerID: string, modelID: string, agent?: string, currentSessionId?: string, attachments?: AttachedFile[], agentMentionName?: string | null) => {
                    if (!currentSessionId) {
                        throw new Error("No session selected");
                    }

                    const sessionId = currentSessionId;

                    if (get().sessionAbortFlags.has(sessionId)) {
                        set((state) => {
                            const nextAbortFlags = new Map(state.sessionAbortFlags);
                            nextAbortFlags.delete(sessionId);
                            return { sessionAbortFlags: nextAbortFlags };
                        });
                    }

                    await executeWithSessionDirectory(sessionId, async () => {
                        try {
                            let effectiveContent = content;
                            const isCommand = content.startsWith("/");

                            if (isCommand) {
                                const spaceIndex = content.indexOf(" ");
                                const command = spaceIndex === -1 ? content.substring(1) : content.substring(1, spaceIndex);
                                const commandArgs = spaceIndex === -1 ? "" : content.substring(spaceIndex + 1).trim();

                                const apiClient = opencodeClient.getApiClient();
                                const directory = opencodeClient.getDirectory();

                                if (command === "init") {
                                    const messageId = `msg_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`;

                                    await apiClient.session.init({
                                        path: { id: sessionId },
                                        body: {
                                            messageID: messageId,
                                            providerID,
                                            modelID,
                                        },
                                        query: directory ? { directory } : undefined,
                                    });

                                    return;
                                }

                                if (command === "summarize") {
                                    await apiClient.session.summarize({
                                        path: { id: sessionId },
                                        body: {
                                            providerID,
                                            modelID,
                                        },
                                        query: directory ? { directory } : undefined,
                                    });

                                    return;
                                }

                                try {
                                    const commandDetails = await opencodeClient.getCommandDetails(command);
                                    if (commandDetails?.template) {
                                        effectiveContent = commandDetails.template.replace(/\$ARGUMENTS/g, commandArgs);
                                    } else {
                                        effectiveContent = content;
                                    }
                                } catch (error) {
                                    console.error("Command template resolution failed:", error);
                                    effectiveContent = content;
                                }
                            }

                            set({
                                lastUsedProvider: { providerID, modelID },
                            });

                            set((state) => {
                                const memoryState = state.sessionMemoryState.get(sessionId) || {
                                    viewportAnchor: 0,
                                    isStreaming: false,
                                    lastAccessedAt: Date.now(),
                                    backgroundMessageCount: 0,
                                };

                                const existingTimer = streamingCooldownTimers.get(sessionId);
                                if (existingTimer) {
                                    clearTimeout(existingTimer);
                                    streamingCooldownTimers.delete(sessionId);
                                }

                                const newMemoryState = new Map(state.sessionMemoryState);
                                newMemoryState.set(sessionId, {
                                    ...memoryState,
                                    isStreaming: true,
                                    streamStartTime: Date.now(),
                                    streamingCooldownUntil: undefined,
                                });
                                return { sessionMemoryState: newMemoryState };
                            });

                            try {
                                const controller = new AbortController();
                                set((state) => {
                                    const nextControllers = new Map(state.abortControllers);
                                    nextControllers.set(sessionId, controller);
                                    return { abortControllers: nextControllers };
                                });

                                const filePayloads = (attachments ?? []).map((file) => ({
                                    type: "file" as const,
                                    mime: file.mimeType,
                                    filename: file.filename,
                                    url: file.dataUrl,
                                }));

                                set((state) => {
                                    const next = new Set(state.pendingAssistantHeaderSessions);
                                    next.add(sessionId);
                                    const nextUserMeta = new Map(state.pendingUserMessageMetaBySession);
                                    nextUserMeta.set(sessionId, {
                                        mode: typeof agent === 'string' && agent.trim().length > 0 ? agent.trim() : undefined,
                                        providerID,
                                        modelID,
                                    });
                                    return { pendingAssistantHeaderSessions: next, pendingUserMessageMetaBySession: nextUserMeta };
                                });

                                await opencodeClient.sendMessage({
                                    id: sessionId,
                                    providerID,
                                    modelID,
                                    text: effectiveContent,
                                    agent,
                                    files: filePayloads.length > 0 ? filePayloads : undefined,
                                    agentMentions: agentMentionName ? [{ name: agentMentionName }] : undefined,
                                });

                                if (filePayloads.length > 0) {
                                    try {
                                        useFileStore.getState().clearAttachedFiles();
                                    } catch (clearError) {
                                        console.error("Failed to clear attached files after send", clearError);
                                    }
                                }
                                set((state) => {
                                    const nextControllers = new Map(state.abortControllers);
                                    nextControllers.delete(sessionId);
                                    return { abortControllers: nextControllers };
                                });
                            } catch (error: any) {
                                let errorMessage = "Network error while sending message. The message may still be processing.";

                                if (error.name === "AbortError") {
                                    errorMessage = "Request timed out. The message may still be processing.";
                                } else if (error.message?.includes("504") || error.message?.includes("Gateway")) {
                                    errorMessage = "Gateway timeout - your message is being processed. Please wait for response.";
                                    set((state) => {
                                        const nextControllers = new Map(state.abortControllers);
                                        nextControllers.delete(sessionId);
                                        return { abortControllers: nextControllers };
                                    });
                                    return;
                                } else if (error.message) {
                                    errorMessage = error.message;
                                }

                                set((state) => {
                                    const nextControllers = new Map(state.abortControllers);
                                    nextControllers.delete(sessionId);
                                    const nextHeaders = new Set(state.pendingAssistantHeaderSessions);
                                    nextHeaders.delete(sessionId);
                                    const nextUserMeta = new Map(state.pendingUserMessageMetaBySession);
                                    nextUserMeta.delete(sessionId);
                                    return { abortControllers: nextControllers, pendingAssistantHeaderSessions: nextHeaders, pendingUserMessageMetaBySession: nextUserMeta };
                                });

                                throw new Error(errorMessage);
                            }
                        } catch (error: any) {
                            let errorMessage = "Network error while sending message. The message may still be processing.";

                            if (error.name === "AbortError") {
                                errorMessage = "Request timed out. The message may still be processing.";
                            } else if (error.response?.status === 401) {
                                errorMessage = "Session not found or unauthorized. Please refresh the page.";
                            } else if (error.response?.status === 502) {
                                errorMessage = "OpenCode is restarting. Please wait a moment and try again.";
                            } else if (error.message?.includes("504") || error.message?.includes("Gateway")) {
                                errorMessage = "Gateway timeout - your message is being processed. Please wait for response.";
                            } else if (error.message) {
                                errorMessage = error.message;
                            }

                            set((state) => {
                                const nextControllers = new Map(state.abortControllers);
                                nextControllers.delete(sessionId);
                                const nextHeaders = new Set(state.pendingAssistantHeaderSessions);
                                nextHeaders.delete(sessionId);
                                const nextUserMeta = new Map(state.pendingUserMessageMetaBySession);
                                nextUserMeta.delete(sessionId);
                                return { abortControllers: nextControllers, pendingAssistantHeaderSessions: nextHeaders, pendingUserMessageMetaBySession: nextUserMeta };
                            });

                            throw new Error(errorMessage);
                        }
                    });
                },

                abortCurrentOperation: async (currentSessionId?: string) => {
                    if (!currentSessionId) {
                        return;
                    }

                    const stateSnapshot = get();
                    const { abortControllers, messages: storeMessages } = stateSnapshot;

                    const controller = abortControllers.get(currentSessionId);
                    controller?.abort();

                    const activeIds = collectActiveMessageIdsForSession(stateSnapshot, currentSessionId);

                    if (activeIds.size === 0) {
                        const sessionMessages = currentSessionId ? storeMessages.get(currentSessionId) ?? [] : [];
                        let fallbackAssistantId: string | null = null;
                        for (let index = sessionMessages.length - 1; index >= 0; index -= 1) {
                            const message = sessionMessages[index];
                            if (!message || message.info.role !== 'assistant') {
                                continue;
                            }

                            if (!fallbackAssistantId) {
                                fallbackAssistantId = message.info.id;
                            }

                            const hasWorkingPart = (message.parts ?? []).some((part) => {
                                return part.type === 'reasoning' || part.type === 'tool' || part.type === 'step-start';
                            });
                            if (hasWorkingPart) {
                                activeIds.add(message.info.id);
                                break;
                            }
                        }

                        if (activeIds.size === 0 && fallbackAssistantId) {
                            activeIds.add(fallbackAssistantId);
                        }
                    }

                    for (const id of activeIds) {
                        const timeout = timeoutRegistry.get(id);
                        if (timeout) {
                            clearTimeout(timeout);
                            timeoutRegistry.delete(id);
                            lastContentRegistry.delete(id);
                        }
                    }

                    if (activeIds.size > 0) {
                        clearLifecycleTimersForIds(activeIds);
                    }

                    const abortTimestamp = Date.now();

                    set((state) => {
                        const updatedStates = removeLifecycleEntries(state.messageStreamStates, activeIds);

                        const sessionMessages = state.messages.get(currentSessionId) ?? [];
                        let messagesChanged = false;
                        let updatedMessages = state.messages;

                        if (sessionMessages.length > 0 && activeIds.size > 0) {
                            const updatedSessionMessages = sessionMessages.map((message) => {
                                if (!activeIds.has(message.info.id) && activeIds.size > 0) {
                                    return message;
                                }

                                const updatedParts = (message.parts ?? []).map((part) => {
                                    if (part.type === 'reasoning') {
                                        const reasoningPart = part as any;
                                        const time = { ...(reasoningPart.time ?? {}) };
                                        if (typeof time.end !== 'number') {
                                            time.end = abortTimestamp;
                                        }
                                        return {
                                            ...reasoningPart,
                                            time,
                                        } as Part;
                                    }

                                    if (part.type === 'tool') {
                                        const toolPart = part as any;
                                        const stateData = { ...(toolPart.state ?? {}) };
                                        if (stateData.status === 'running' || stateData.status === 'pending') {
                                            stateData.status = 'aborted';
                                        }
                                        return {
                                            ...toolPart,
                                            state: stateData,
                                        } as Part;
                                    }

                                    if (part.type === 'step-start') {
                                        const stepPart = part as any;
                                        return {
                                            ...stepPart,
                                            type: 'step-finish',
                                            aborted: true,
                                        } as Part;
                                    }

                                    return part;
                                });

                                messagesChanged = true;
                                return {
                                    ...message,
                                    info: {
                                        ...message.info,
                                        abortedAt: abortTimestamp,
                                        streaming: false,
                                        status: 'aborted',
                                    },
                                    parts: updatedParts,
                                };
                            });

                            if (messagesChanged) {
                                updatedMessages = new Map(state.messages);
                                updatedMessages.set(currentSessionId, updatedSessionMessages);
                            }
                        }
                        const memoryState = state.sessionMemoryState.get(currentSessionId);
                        let nextMemoryState = state.sessionMemoryState;
                        if (memoryState) {
                            const updatedMemory = new Map(state.sessionMemoryState);
                            updatedMemory.set(currentSessionId, {
                                ...memoryState,
                                isStreaming: false,
                                streamStartTime: undefined,
                                isZombie: false,
                            });
                            nextMemoryState = updatedMemory;
                        }

                        const nextAbortFlags = new Map(state.sessionAbortFlags);
                        nextAbortFlags.set(currentSessionId, {
                            timestamp: abortTimestamp,
                            acknowledged: false,
                        });

                        return {
                            messageStreamStates: updatedStates,
                            sessionMemoryState: nextMemoryState,
                            sessionAbortFlags: nextAbortFlags,
                            abortControllers: (() => {
                                const nextControllers = new Map(state.abortControllers);
                                nextControllers.delete(currentSessionId);
                                return nextControllers;
                            })(),
                            streamingMessageIds: setStreamingIdForSession(state.streamingMessageIds, currentSessionId, null),
                            ...(messagesChanged ? { messages: updatedMessages } : {}),
                        };
                    });

                    void opencodeClient.abortSession(currentSessionId).catch((error) => {
                        console.warn('Abort request failed:', error);
                    });
                },

                _addStreamingPartImmediate: (sessionId: string, messageId: string, part: Part, role?: string, currentSessionId?: string) => {
                    const stateSnapshot = get();
                    if (ignoredAssistantMessageIds.has(messageId)) {
                        return;
                    }

                    const trimmedHeadMaxId = stateSnapshot.sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;
                    if (trimmedHeadMaxId && !isIdNewer(messageId, trimmedHeadMaxId)) {
                        (window as any).__messageTracker?.(messageId, 'ignored_trimmed_stream_part');
                        return;
                    }

                    const existingMessagesSnapshot = stateSnapshot.messages.get(sessionId) || [];
                    const existingMessageSnapshot = existingMessagesSnapshot.find((m) => m.info.id === messageId);

                    const actualRole = (() => {
                        if (role === 'user') return 'user';
                        if (existingMessageSnapshot?.info.role === 'user') return 'user';
                        return role || existingMessageSnapshot?.info.role || 'assistant';
                    })();

                    const memoryStateSnapshot = get().sessionMemoryState.get(sessionId);
                    if (memoryStateSnapshot?.streamStartTime) {
                        const streamDuration = Date.now() - memoryStateSnapshot.streamStartTime;
                        if (streamDuration > MEMORY_LIMITS.ZOMBIE_TIMEOUT) {
                            if (!memoryStateSnapshot.isZombie) {
                                set((state) => {
                                    const newMemoryState = new Map(state.sessionMemoryState);
                                    newMemoryState.set(sessionId, {
                                        ...memoryStateSnapshot,
                                        isZombie: true,
                                    });
                                    return { sessionMemoryState: newMemoryState };
                                });
                            }

                            setTimeout(() => {
                                const store = get();
                                store.completeStreamingMessage(sessionId, messageId);
                            }, 0);
                            (window as any).__messageTracker?.(messageId, 'skipped_zombie_stream');
                            return;
                        }
                    }

                    set((state) => {
                        const sessionMessages = state.messages.get(sessionId) || [];
                        const messagesArray = [...sessionMessages];
                        const updates: any = {};

                        const indexedSessions = upsertMessageSessionIndex(state.messageSessionIndex, messageId, sessionId);
                        if (indexedSessions !== state.messageSessionIndex) {
                            updates.messageSessionIndex = indexedSessions;
                        }

                        const finalizeAbortState = (result: Partial<MessageState>): Partial<MessageState> => {
                            const shouldClearAbortFlag =
                                (actualRole === 'assistant' || actualRole === 'user') &&
                                state.sessionAbortFlags.has(sessionId);
                            if (!shouldClearAbortFlag) {
                                return result;
                            }
                            const nextAbortFlags = new Map(state.sessionAbortFlags);
                            nextAbortFlags.delete(sessionId);
                            return {
                                ...result,
                                sessionAbortFlags: nextAbortFlags,
                            };
                        };

                        const maintainTimeouts = (text: string) => {
                            const value = text || '';
                            const lastContent = lastContentRegistry.get(messageId);

                            if (value && lastContent === value) {
                                const currentState = get();
                                if (isMessageStreamingInSession(currentState, sessionId, messageId)) {
                                    const existingTimeout = timeoutRegistry.get(messageId);
                                    if (existingTimeout) {
                                        clearTimeout(existingTimeout);
                                        timeoutRegistry.delete(messageId);
                                    }
                                    setTimeout(() => {
                                        const store = get();
                                        if (typeof store.forceCompleteMessage === "function") {
                                            store.forceCompleteMessage(sessionId, messageId, "timeout");
                                        }
                                        store.completeStreamingMessage(sessionId, messageId);
                                    }, 100);
                                }
                            }

                            lastContentRegistry.set(messageId, value);

                            const existingTimeout = timeoutRegistry.get(messageId);
                            if (existingTimeout) {
                                clearTimeout(existingTimeout);
                            }
                            const newTimeout = setTimeout(() => {
                                const store = get();
                                if (typeof store.forceCompleteMessage === "function") {
                                    store.forceCompleteMessage(sessionId, messageId, "timeout");
                                }
                                if (isMessageStreamingInSession(store, sessionId, messageId)) {
                                    store.completeStreamingMessage(sessionId, messageId);
                                }
                                timeoutRegistry.delete(messageId);
                                lastContentRegistry.delete(messageId);
                            }, 8000);
                            timeoutRegistry.set(messageId, newTimeout);
                        };

                        const isBackgroundSession = sessionId !== currentSessionId;
                        const memoryState = state.sessionMemoryState.get(sessionId);
                        if (isBackgroundSession && memoryState?.isStreaming) {
                            if (messagesArray.length >= MEMORY_LIMITS.BACKGROUND_STREAMING_BUFFER) {
                                messagesArray.shift();
                            }

                            const newMemoryState = new Map(state.sessionMemoryState);
                            newMemoryState.set(sessionId, {
                                ...memoryState,
                                backgroundMessageCount: (memoryState.backgroundMessageCount || 0) + 1,
                            });
                            state.sessionMemoryState = newMemoryState;
                        }

                        if (actualRole === 'assistant') {
                            const currentMemoryState = state.sessionMemoryState.get(sessionId);
                            if (currentMemoryState) {
                                const now = Date.now();
                                const nextMemoryState = new Map(state.sessionMemoryState);
                                nextMemoryState.set(sessionId, {
                                    ...currentMemoryState,
                                    isStreaming: true,
                                    streamStartTime: currentMemoryState.streamStartTime ?? now,
                                    lastAccessedAt: now,
                                    isZombie: false,
                                });
                                state.sessionMemoryState = nextMemoryState;
                            }
                        }

                    const incomingText = extractTextFromPart(part);
                    if (isExecutionForkMetaText(incomingText)) {
                        (part as any).synthetic = true;
                    }
                    if (streamDebugEnabled() && actualRole === "assistant") {
                        try {
                            console.info("[STREAM-TRACE] part", {
                                messageId,
                                role: actualRole,
                                type: (part as any)?.type || "text",
                                textLen: incomingText.length,
                                snapshotParts: existingMessagesSnapshot.length,
                            });
                        } catch { /* ignored */ }
                    }

                        const previousStreamingMap = updates.streamingMessageIds ?? state.streamingMessageIds;
                        if (actualRole === 'assistant') {
                            const nextStreamingMap = setStreamingIdForSession(previousStreamingMap, sessionId, messageId);
                            if (nextStreamingMap !== previousStreamingMap) {
                                updates.streamingMessageIds = nextStreamingMap;
                                (window as any).__messageTracker?.(messageId, 'streamingId_set_latest');
                            }
                        }

                        const messageIndex = messagesArray.findIndex((m) => m.info.id === messageId);

                        if (messageIndex !== -1 && actualRole === 'user') {
                            const existingMessage = messagesArray[messageIndex];
                            const existingPartIndex = existingMessage.parts.findIndex((p) => p.id === part.id);

                            if ((part as any).synthetic === true) {
                                (window as any).__messageTracker?.(messageId, 'skipped_synthetic_user_part');
                                return state;
                            }

                            const normalizedPart = normalizeStreamingPart(
                                part,
                                existingPartIndex !== -1 ? existingMessage.parts[existingPartIndex] : undefined
                            );
                            (window as any).__messageTracker?.(messageId, `user_part_type:${(normalizedPart as any).type || 'unknown'}`);

                            const updatedMessage = { ...existingMessage };
                            if (existingPartIndex !== -1) {
                                updatedMessage.parts = updatedMessage.parts.map((p, idx) =>
                                    idx === existingPartIndex ? normalizedPart : p
                                );
                            } else {
                                updatedMessage.parts = [...updatedMessage.parts, normalizedPart];
                            }

                            const updatedMessages = [...messagesArray];
                            updatedMessages[messageIndex] = updatedMessage;

                            const newMessages = new Map(state.messages);
                            newMessages.set(sessionId, updatedMessages);

                            return finalizeAbortState({ messages: newMessages });
                        }

                        if (actualRole === 'assistant' && messageIndex !== -1) {

                            const existingMessage = messagesArray[messageIndex];
                            const existingPartIndex = existingMessage.parts.findIndex((p) => p.id === part.id);

                            const normalizedPart = normalizeStreamingPart(
                                part,
                                existingPartIndex !== -1 ? existingMessage.parts[existingPartIndex] : undefined
                            );
                            (window as any).__messageTracker?.(messageId, `part_type:${(normalizedPart as any).type || 'unknown'}`);

                            const updatedMessage = { ...existingMessage };
                            if (existingPartIndex !== -1) {
                                updatedMessage.parts = updatedMessage.parts.map((p, idx) =>
                                    idx === existingPartIndex ? normalizedPart : p
                                );
                            } else {
                                updatedMessage.parts = [...updatedMessage.parts, normalizedPart];
                            }

                            const updatedMessages = [...messagesArray];
                            updatedMessages[messageIndex] = updatedMessage;

                            const newMessages = new Map(state.messages);
                            newMessages.set(sessionId, updatedMessages);

                            updates.messageStreamStates = touchStreamingLifecycle(state.messageStreamStates, messageId);
                            const nextStreamingMap = setStreamingIdForSession(updates.streamingMessageIds ?? state.streamingMessageIds, sessionId, messageId);
                            if (nextStreamingMap !== (updates.streamingMessageIds ?? state.streamingMessageIds)) {
                                updates.streamingMessageIds = nextStreamingMap;
                                (window as any).__messageTracker?.(messageId, 'streamingId_set');
                            }

                            if ((normalizedPart as any).type === 'text') {
                                maintainTimeouts((normalizedPart as any).text || '');
                            } else {
                                maintainTimeouts('');
                            }

                            return finalizeAbortState({ messages: newMessages, ...updates });
                        }

                        if (messageIndex === -1) {

                            if (actualRole === 'user') {

                                if ((part as any).synthetic === true) {
                                    (window as any).__messageTracker?.(messageId, 'skipped_synthetic_new_user_part');
                                    return state;
                                }

                                const normalizedPart = normalizeStreamingPart(part);
                                (window as any).__messageTracker?.(messageId, `new_user_part_type:${(normalizedPart as any).type || 'unknown'}`);

                                const pendingMeta = state.pendingUserMessageMetaBySession.get(sessionId);
                                const contextStore = useContextStore.getState();
                                const sessionAgent =
                                    pendingMeta?.mode ??
                                    contextStore.getSessionAgentSelection(sessionId) ??
                                    contextStore.getCurrentAgent(sessionId);
                                const agentMode = typeof sessionAgent === 'string' && sessionAgent.trim().length > 0
                                    ? sessionAgent.trim()
                                    : undefined;
                                const providerID = pendingMeta?.providerID ?? (state.lastUsedProvider?.providerID || undefined);
                                const modelID = pendingMeta?.modelID ?? (state.lastUsedProvider?.modelID || undefined);

                                if (pendingMeta) {
                                    updates.pendingUserMessageMetaBySession = cleanupPendingUserMessageMeta(state.pendingUserMessageMetaBySession, sessionId);
                                }

                                const newUserMessage = {
                                    info: {
                                        id: messageId,
                                        sessionID: sessionId,
                                        role: 'user' as const,
                                        clientRole: 'user',
                                        userMessageMarker: true,
                                        ...(agentMode ? { mode: agentMode } : {}),
                                        ...(providerID ? { providerID } : {}),
                                        ...(modelID ? { modelID } : {}),
                                        time: {
                                            created: Date.now(),
                                        },
                                    },
                                    parts: [normalizedPart],
                                };

                                const updatedMessages = [...messagesArray, newUserMessage];

                                updatedMessages.sort((a, b) => {
                                    const aTime = (a.info as any)?.time?.created || 0;
                                    const bTime = (b.info as any)?.time?.created || 0;
                                    return aTime - bTime;
                                });

                                const newMessages = new Map(state.messages);
                                newMessages.set(sessionId, updatedMessages);

                                return finalizeAbortState({ messages: newMessages });
                            }

                            if ((part as any)?.type === 'text') {
                                const textIncoming = extractTextFromPart(part).trim();
                                if (textIncoming.length > 0) {
                                    const latestUser = [...messagesArray]
                                        .reverse()
                                        .find((m) => m.info.role === 'user');
                                    if (latestUser) {
                                        const latestUserText = latestUser.parts.map((p) => extractTextFromPart(p)).join('').trim();
                                        if (latestUserText.length > 0 && latestUserText === textIncoming) {
                                            ignoredAssistantMessageIds.add(messageId);
                                            (window as any).__messageTracker?.(messageId, 'ignored_assistant_echo');
                                            return state;
                                        }
                                    }
                                }
                            }

                            const normalizedPart = normalizeStreamingPart(part);
                            (window as any).__messageTracker?.(messageId, `part_type:${(normalizedPart as any).type || 'unknown'}`);

                            if ((normalizedPart as any).type === 'text') {
                                maintainTimeouts((normalizedPart as any).text || '');
                            } else {
                                maintainTimeouts('');
                            }

                            const pendingEntry = state.pendingAssistantParts.get(messageId);
                            const pendingParts = pendingEntry ? [...pendingEntry.parts] : [];
                            const pendingIndex = pendingParts.findIndex((existing) => existing.id === normalizedPart.id);

                            if (pendingIndex !== -1) {
                                pendingParts[pendingIndex] = normalizedPart;
                            } else {
                                pendingParts.push(normalizedPart);
                            }

                            const newPending = new Map(state.pendingAssistantParts);
                            newPending.set(messageId, { sessionId, parts: pendingParts });

                            const providerID = state.lastUsedProvider?.providerID || "";
                            const modelID = state.lastUsedProvider?.modelID || "";
                            const now = Date.now();
                            const cwd = opencodeClient.getDirectory() ?? "/";
                            const contextStore = useContextStore.getState();
                            const sessionAgent = contextStore.getSessionAgentSelection(sessionId)
                                ?? contextStore.getCurrentAgent(sessionId);
                            const agentMode = typeof sessionAgent === "string" && sessionAgent.trim().length > 0
                                ? sessionAgent.trim()
                                : undefined;

                            const shouldAnchorHeader = state.pendingAssistantHeaderSessions.has(sessionId);
                            if (shouldAnchorHeader) {
                                const nextPendingHeaders = new Set(state.pendingAssistantHeaderSessions);
                                nextPendingHeaders.delete(sessionId);
                                updates.pendingAssistantHeaderSessions = nextPendingHeaders;
                            }

                            const placeholderInfo = (actualRole === "user"
                                ? {
                                    id: messageId,
                                    sessionID: sessionId,
                                    role: "user",
                                    time: { created: now },
                                    agent: agentMode || "default",
                                    model: { providerID, modelID },
                                    clientRole: actualRole,
                                    animationSettled: undefined,
                                    streaming: undefined,
                                }
                                : {
                                    id: messageId,
                                    sessionID: sessionId,
                                    role: "assistant",
                                    time: { created: now },
                                    parentID: messageId,
                                    modelID,
                                    providerID,
                                    mode: agentMode || "default",
                                    ...(shouldAnchorHeader ? { openchamberHeaderAnchor: true } : {}),
                                    path: { cwd, root: cwd },
                                    cost: 0,
                                    tokens: {
                                        input: 0,
                                        output: 0,
                                        reasoning: 0,
                                        cache: { read: 0, write: 0 },
                                    },
                                    clientRole: actualRole,
                                    animationSettled: false,
                                    streaming: true,
                                }) as unknown as Message;

                            const placeholderMessage = {
                                info: placeholderInfo,
                                parts: pendingParts,
                            };

                            const nextMessages = [...messagesArray, placeholderMessage];

                            const newMessages = new Map(state.messages);
                            newMessages.set(sessionId, nextMessages);

                            if (actualRole === 'assistant') {
                                updates.messageStreamStates = touchStreamingLifecycle(state.messageStreamStates, messageId);

                                const nextStreamingMap = setStreamingIdForSession(updates.streamingMessageIds ?? state.streamingMessageIds, sessionId, messageId);
                                if (nextStreamingMap !== (updates.streamingMessageIds ?? state.streamingMessageIds)) {
                                    updates.streamingMessageIds = nextStreamingMap;
                                    (window as any).__messageTracker?.(messageId, 'streamingId_set');
                                }
                            }

                            return finalizeAbortState({
                                messages: newMessages,
                                pendingAssistantParts: newPending,
                                ...updates,
                            });
                        } else {

                            const existingMessage = messagesArray[messageIndex];
                            const existingPartIndex = existingMessage.parts.findIndex((p) => p.id === part.id);

                            const normalizedPart = normalizeStreamingPart(
                                part,
                                existingPartIndex !== -1 ? existingMessage.parts[existingPartIndex] : undefined
                            );
                            (window as any).__messageTracker?.(messageId, `part_type:${(normalizedPart as any).type || 'unknown'}`);

                            const updatedMessage = { ...existingMessage };
                            if (existingPartIndex !== -1) {
                                updatedMessage.parts = updatedMessage.parts.map((p, idx) =>
                                    idx === existingPartIndex ? normalizedPart : p
                                );
                            } else {
                                updatedMessage.parts = [...updatedMessage.parts, normalizedPart];
                            }

                            const updatedMessages = [...messagesArray];
                            updatedMessages[messageIndex] = updatedMessage;

                            const newMessages = new Map(state.messages);
                            newMessages.set(sessionId, updatedMessages);

                            if (updatedMessage.info.role === "assistant") {
                                updates.messageStreamStates = touchStreamingLifecycle(state.messageStreamStates, messageId);
                                const nextStreamingMap = setStreamingIdForSession(updates.streamingMessageIds ?? state.streamingMessageIds, sessionId, messageId);
                                if (nextStreamingMap !== (updates.streamingMessageIds ?? state.streamingMessageIds)) {
                                    updates.streamingMessageIds = nextStreamingMap;
                                    (window as any).__messageTracker?.(messageId, 'streamingId_set');
                                }
                            }

                            if ((normalizedPart as any).type === 'text') {
                                maintainTimeouts((normalizedPart as any).text || '');
                            } else {
                                maintainTimeouts('');
                            }

                            return finalizeAbortState({ messages: newMessages, ...updates });
                        }
                    });

                    const partType = (part as any)?.type;
                    if (partType === 'step-finish' && actualRole !== 'user') {
                        setTimeout(() => {
                            const store = get();
                            store.completeStreamingMessage(sessionId, messageId);
                        }, 0);
                    }
                },

                addStreamingPart: (sessionId: string, messageId: string, part: Part, role?: string, currentSessionId?: string) => {

                    if (role !== 'user') {
                        get()._addStreamingPartImmediate(sessionId, messageId, part, role, currentSessionId);
                        return;
                    }

                    batchQueue.push({ sessionId, messageId, part, role, currentSessionId });

                    if (!flushTimer) {
                        flushTimer = setTimeout(() => {
                            const itemsToProcess = [...batchQueue];
                            batchQueue = [];
                            flushTimer = null;

                            const store = get();
                            for (const item of itemsToProcess) {
                                store._addStreamingPartImmediate(item.sessionId, item.messageId, item.part, item.role, item.currentSessionId);
                            }
                        }, USER_BATCH_WINDOW_MS);
                    }
                },

                forceCompleteMessage: (sessionId: string | null | undefined, messageId: string, source: "timeout" | "cooldown" = "timeout") => {
                    const resolveSessionId = (state: MessageState): string | null => {
                        if (sessionId) {
                            return sessionId;
                        }
                        for (const [candidateId, sessionMessages] of state.messages.entries()) {
                            if (sessionMessages.some((msg) => msg.info.id === messageId)) {
                                return candidateId;
                            }
                        }
                        return null;
                    };

                    set((state) => {
                        const targetSessionId = resolveSessionId(state);
                        if (!targetSessionId) {
                            return state;
                        }

                        const sessionMessages = state.messages.get(targetSessionId) ?? [];
                        const messageIndex = sessionMessages.findIndex((msg) => msg.info.id === messageId);
                        if (messageIndex === -1) {
                            return state;
                        }

                        const message = sessionMessages[messageIndex];
                        if (!message) {
                            return state;
                        }

                        const now = Date.now();
                        const existingInfo = message.info as any;
                        const existingCompleted = typeof existingInfo?.time?.completed === "number" && existingInfo.time.completed > 0;

                        let infoChanged = false;
                        const updatedInfo: Record<string, any> = { ...existingInfo };

                        if (!existingCompleted) {
                            updatedInfo.time = {
                                ...(existingInfo.time ?? {}),
                                completed: now,
                            };
                            infoChanged = true;
                        }

                        if (updatedInfo.status !== "completed") {
                            updatedInfo.status = "completed";
                            infoChanged = true;
                        }

                        if (updatedInfo.streaming) {
                            updatedInfo.streaming = false;
                            infoChanged = true;
                        }

                        let partsChanged = false;
                        const updatedParts = message.parts.map((part) => {
                            if (!part) {
                                return part;
                            }

                            if (part.type === "tool") {
                                const existingState = (part as any).state;
                                if (!existingState) {
                                    return part;
                                }

                                const status = existingState.status;
                                const needsStatusUpdate = status === "running" || status === "pending" || status === "started";
                                const needsEndTimestamp = !existingState.time || typeof existingState.time?.end !== "number";

                                if (needsStatusUpdate || needsEndTimestamp) {
                                    const nextState: Record<string, any> = { ...existingState };
                                    if (needsStatusUpdate) {
                                        nextState.status = "completed";
                                    }
                                    if (needsEndTimestamp) {
                                        nextState.time = {
                                            ...(existingState.time ?? {}),
                                            end: now,
                                        };
                                    }
                                    partsChanged = true;
                                    return {
                                        ...part,
                                        state: nextState,
                                    } as Part;
                                }
                                return part;
                            }

                            if (part.type === "reasoning") {
                                const reasoningTime = (part as any).time;
                                if (!reasoningTime || typeof reasoningTime.end !== "number") {
                                    partsChanged = true;
                                    return {
                                        ...part,
                                        time: {
                                            ...(reasoningTime ?? {}),
                                            end: now,
                                        },
                                    } as Part;
                                }
                                return part;
                            }

                            if (part.type === "text") {
                                const textTime = (part as any).time;
                                if (textTime && typeof textTime.end !== "number") {
                                    partsChanged = true;
                                    return {
                                        ...part,
                                        time: {
                                            ...textTime,
                                            end: now,
                                        },
                                    } as Part;
                                }
                                return part;
                            }

                            return part;
                        });

                        if (!infoChanged && !partsChanged) {
                            return state;
                        }

                        (window as any).__messageTracker?.(messageId, `force_complete:${source}`);

                        const updatedMessage = {
                            ...message,
                            info: updatedInfo as Message,
                            parts: partsChanged ? updatedParts : message.parts,
                        };

                        const nextSessionMessages = [...sessionMessages];
                        nextSessionMessages[messageIndex] = updatedMessage;

                        const nextMessages = new Map(state.messages);
                        nextMessages.set(targetSessionId, nextSessionMessages);

                        return { messages: nextMessages };
                    });
                },

                markMessageStreamSettled: (messageId: string) => {
                    set((state) => {

                        clearLifecycleCompletionTimer(messageId);
                        const next = new Map(state.messageStreamStates);
                        next.delete(messageId);

                        let updatedMessages = state.messages;
                        let messagesModified = false;

                        state.messages.forEach((sessionMessages, sessionId) => {
                            if (messagesModified) return;
                            const idx = sessionMessages.findIndex((msg) => msg.info.id === messageId);
                            if (idx === -1) return;

                            const message = sessionMessages[idx];
                            if ((message.info as any)?.animationSettled) {
                                return;
                            }

                            const updatedMessage = {
                                ...message,
                                info: {
                                    ...message.info,
                                    animationSettled: true,
                                },
                            };

                            const sessionArray = [...sessionMessages];
                            sessionArray[idx] = updatedMessage;

                            const newMessages = new Map(state.messages);
                            newMessages.set(sessionId, sessionArray);
                            updatedMessages = newMessages;
                            messagesModified = true;
                        });

                        const updates: Partial<MessageState> & { messageStreamStates: Map<string, MessageStreamLifecycle> } = {
                            messageStreamStates: next,
                            ...(messagesModified ? { messages: updatedMessages } : {}),
                        } as any;

                        return updates;
                    });
                    clearLifecycleTimersForIds([messageId]);
                },

                updateMessageInfo: (sessionId: string, messageId: string, messageInfo: any) => {
                    set((state) => {
                        const trimmedHeadMaxId = state.sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;
                        if (trimmedHeadMaxId && !isIdNewer(messageId, trimmedHeadMaxId)) {
                            (window as any).__messageTracker?.(messageId, 'ignored_trimmed_update');
                            return state;
                        }

                        const sessionMessages = state.messages.get(sessionId) ?? [];
                        const normalizedSessionMessages = [...sessionMessages];

                        const messageIndex = normalizedSessionMessages.findIndex((msg) => msg.info.id === messageId);
                        const pendingEntry = state.pendingAssistantParts.get(messageId);

                        const mergeParts = (existingParts: Part[] = [], incomingParts: Part[] = []) => {
                            if (!incomingParts.length) {
                                return existingParts;
                            }
                            const merged = [...existingParts];
                            incomingParts.forEach((incomingPart) => {
                                const idx = merged.findIndex((part) => part.id === incomingPart.id);
                                if (idx === -1) {
                                    merged.push(incomingPart);
                                } else {
                                    merged[idx] = incomingPart;
                                }
                            });
                            return merged;
                        };

                        const ensureClientRole = (info: any) => {

                            if (!info) {
                                return info;
                            }
                            const clientRole = info.clientRole ?? info.role;
                            const userMarker = clientRole === 'user' ? true : info.userMessageMarker;
                            return {
                                ...info,
                                clientRole,
                                ...(userMarker ? { userMessageMarker: true } : {}),
                            };
                        };

                        if (messageIndex === -1) {
                            console.info("[MESSAGE-DEBUG] updateMessageInfo: messageIndex === -1", {
                                sessionId,
                                messageId,
                                messageInfo,
                                existingCount: normalizedSessionMessages.length,
                            });

                            if (normalizedSessionMessages.length > 0) {
                                const firstMessage = normalizedSessionMessages[0];
                                const firstInfo = firstMessage?.info as any;
                                const firstCreated = typeof firstInfo?.time?.created === 'number' ? firstInfo.time.created : null;
                                const firstId = typeof firstInfo?.id === 'string' ? firstInfo.id : null;

                                const incomingInfoToCompare = messageInfo as any;
                                const incomingCreated = typeof incomingInfoToCompare?.time?.created === 'number'
                                    ? incomingInfoToCompare.time.created
                                    : null;
                                const incomingId = typeof incomingInfoToCompare?.id === 'string' ? incomingInfoToCompare.id : messageId;

                                let isOlderThanViewport = false;
                                if (incomingCreated !== null && firstCreated !== null) {
                                    isOlderThanViewport = incomingCreated < firstCreated;
                                }
                                if (!isOlderThanViewport && incomingId && firstId) {
                                    isOlderThanViewport = incomingId.localeCompare(firstId) < 0;
                                }

                                if (isOlderThanViewport) {
                                    (window as any).__messageTracker?.(messageId, 'skipped_evicted_message_update');
                                    return state;
                                }
                            }

                            const incomingInfo = ensureClientRole(messageInfo);

                            if (incomingInfo && incomingInfo.role === 'user') {
                                const pendingParts = pendingEntry?.parts ?? [];
                                const pendingMeta = state.pendingUserMessageMetaBySession.get(sessionId);
                                const newUserMessage = {
                                    info: {
                                        ...incomingInfo,
                                        userMessageMarker: true,
                                        clientRole: 'user',
                                        ...(pendingMeta?.mode ? { mode: pendingMeta.mode } : {}),
                                        ...(pendingMeta?.providerID ? { providerID: pendingMeta.providerID } : {}),
                                        ...(pendingMeta?.modelID ? { modelID: pendingMeta.modelID } : {}),
                                    } as Message,
                                    parts: pendingParts.length > 0 ? [...pendingParts] : [],
                                };

                                const newMessages = new Map(state.messages);

                                const appended = [...normalizedSessionMessages, newUserMessage];

                                appended.sort((a, b) => {
                                    const aTime = (a.info as any)?.time?.created || 0;
                                    const bTime = (b.info as any)?.time?.created || 0;
                                    return aTime - bTime;
                                });
                                newMessages.set(sessionId, appended);

                                const updates: Partial<MessageState> = {
                                    messages: newMessages,
                              ...(pendingMeta
                                 ? {
                                      pendingUserMessageMetaBySession: cleanupPendingUserMessageMeta(state.pendingUserMessageMetaBySession, sessionId),
                                  }
                                 : {}),
                                };

                                const nextIndex = upsertMessageSessionIndex(
                                    updates.messageSessionIndex ?? state.messageSessionIndex,
                                    messageId,
                                    sessionId
                                );
                                if (nextIndex !== (updates.messageSessionIndex ?? state.messageSessionIndex)) {
                                    updates.messageSessionIndex = nextIndex;
                                }

                                if (pendingEntry) {
                                    const newPending = new Map(state.pendingAssistantParts);
                                    newPending.delete(messageId);
                                    updates.pendingAssistantParts = newPending;
                                }

                                return updates;
                            }

                            if (!incomingInfo || incomingInfo.role !== 'assistant') {
                                return state;
                            }

                            const pendingParts = pendingEntry?.parts ?? [];

                            const shouldAnchorHeader = state.pendingAssistantHeaderSessions.has(sessionId);
                            const newMessage = {
                                info: {
                                    ...incomingInfo,
                                    animationSettled: (incomingInfo as any)?.animationSettled ?? false,
                                    ...(shouldAnchorHeader ? { openchamberHeaderAnchor: true } : {}),
                                } as Message,
                                parts: pendingParts.length > 0 ? [...pendingParts] : [],
                            };

                            const newMessages = new Map(state.messages);

                            const appended = [...normalizedSessionMessages, newMessage];
                            newMessages.set(sessionId, appended);

                            const updates: Partial<MessageState> = {
                                messages: newMessages,
                                ...(shouldAnchorHeader
                                    ? {
                                        pendingAssistantHeaderSessions: (() => {
                                            const nextPendingHeaders = new Set(state.pendingAssistantHeaderSessions);
                                            nextPendingHeaders.delete(sessionId);
                                            return nextPendingHeaders;
                                        })(),
                                    }
                                    : {}),
                            };

                            const nextIndex = upsertMessageSessionIndex(
                                updates.messageSessionIndex ?? state.messageSessionIndex,
                                messageId,
                                sessionId
                            );
                            if (nextIndex !== (updates.messageSessionIndex ?? state.messageSessionIndex)) {
                                updates.messageSessionIndex = nextIndex;
                            }

                            if (pendingEntry) {
                                const newPending = new Map(state.pendingAssistantParts);
                                newPending.delete(messageId);
                                updates.pendingAssistantParts = newPending;
                            }

                            return updates;
                        }

                        const existingMessage = normalizedSessionMessages[messageIndex];

                        const existingInfo = existingMessage.info as any;
                        const isUserMessage =
                            existingInfo.userMessageMarker === true ||
                            existingInfo.clientRole === 'user' ||
                            existingInfo.role === 'user';

                         if (isUserMessage) {
 
                             const updatedInfo = {
                                 ...existingMessage.info,
                                 ...messageInfo,
 
                                 role: 'user',
                                 clientRole: 'user',
                                 userMessageMarker: true,
 
                                 providerID: existingInfo.providerID || undefined,
                                 modelID: existingInfo.modelID || undefined,
                             } as any;

                             const pendingMeta = state.pendingUserMessageMetaBySession.get(sessionId);
                             if (pendingMeta && !updatedInfo.mode && pendingMeta.mode) {
                                 updatedInfo.mode = pendingMeta.mode;
                             }
 
                             const updatedMessage = {
                                 ...existingMessage,
                                 info: updatedInfo
                             };
 
                             const newMessages = new Map(state.messages);
                             const updatedSessionMessages = [...normalizedSessionMessages];
                             updatedSessionMessages[messageIndex] = updatedMessage;
                             newMessages.set(sessionId, updatedSessionMessages);

                             if (pendingMeta) {
                                 const nextPending = new Map(state.pendingUserMessageMetaBySession);
                                 nextPending.delete(sessionId);
                                 return { messages: newMessages, pendingUserMessageMetaBySession: nextPending };
                             }
 
                             return { messages: newMessages };
                         }


                        const updatedInfo = {
                            ...existingMessage.info,
                            ...messageInfo,
                        } as any;

                        if (messageInfo.role && messageInfo.role !== existingMessage.info.role) {
                            updatedInfo.role = existingMessage.info.role;
                        }

                        updatedInfo.clientRole = updatedInfo.clientRole ?? existingMessage.info.clientRole ?? existingMessage.info.role;
                        if (updatedInfo.clientRole === "user") {
                            updatedInfo.userMessageMarker = true;
                        }

                        const mergedParts = pendingEntry?.parts
                            ? mergeParts(existingMessage.parts, pendingEntry.parts)
                            : existingMessage.parts;

                        const updatedMessage = {
                            ...existingMessage,
                            info: updatedInfo,
                            parts: mergedParts,
                        };

                        const newMessages = new Map(state.messages);
                        const updatedSessionMessages = [...normalizedSessionMessages];
                        updatedSessionMessages[messageIndex] = updatedMessage;
                        newMessages.set(sessionId, updatedSessionMessages);

                        const updates: Partial<MessageState> = {
                            messages: newMessages,
                        };

                        if (pendingEntry) {
                            const newPending = new Map(state.pendingAssistantParts);
                            newPending.delete(messageId);
                            updates.pendingAssistantParts = newPending;
                        }

                        return updates;
                    });
                },

                completeStreamingMessage: (sessionId: string, messageId: string) => {
                    const state = get();

                    (window as any).__messageTracker?.(
                        messageId,
                        `completion_called_current:${state.streamingMessageIds.get(sessionId) ?? 'none'}`
                    );

                    if (typeof state.forceCompleteMessage === "function") {
                        state.forceCompleteMessage(sessionId, messageId, "cooldown");
                    }

                    const shouldClearStreamingId = state.streamingMessageIds.get(sessionId) === messageId;
                    if (shouldClearStreamingId) {
                        (window as any).__messageTracker?.(messageId, 'streamingId_cleared');
                    } else {
                        (window as any).__messageTracker?.(messageId, 'streamingId_NOT_cleared_different_id');
                    }

                    const updates: Record<string, any> = {};
                    if (shouldClearStreamingId) {
                        updates.streamingMessageIds = setStreamingIdForSession(state.streamingMessageIds, sessionId, null);
                        updates.abortControllers = (() => {
                            const next = new Map(state.abortControllers);
                            next.delete(sessionId);
                            return next;
                        })();
                    }

                    if (state.messageStreamStates.has(messageId)) {
                        const next = new Map(state.messageStreamStates);
                        next.delete(messageId);
                        updates.messageStreamStates = next;
                    }

                    if (Object.keys(updates).length > 0) {
                        set(updates);
                    }

                    if (state.pendingAssistantParts.has(messageId)) {
                        set((currentState) => {
                            if (!currentState.pendingAssistantParts.has(messageId)) {
                                return currentState;
                            }
                            const nextPending = new Map(currentState.pendingAssistantParts);
                            nextPending.delete(messageId);
                            return { pendingAssistantParts: nextPending };
                        });
                    }

                    clearLifecycleTimersForIds([messageId]);

                    let startedCooldown = false;
                    set((state) => {
                        const memoryState = state.sessionMemoryState.get(sessionId);
                        if (!memoryState || !memoryState.isStreaming) return state;

                        const newMemoryState = new Map(state.sessionMemoryState);
                        const now = Date.now();
                        const updatedMemory: SessionMemoryState = {
                            ...memoryState,
                            isStreaming: false,
                            streamStartTime: undefined,
                            isZombie: false,
                            lastAccessedAt: now,
                            streamingCooldownUntil: now + 2000,
                        };
                        newMemoryState.set(sessionId, updatedMemory);
                        startedCooldown = true;
                        return { sessionMemoryState: newMemoryState };
                    });

                    if (startedCooldown) {
                        const existingTimer = streamingCooldownTimers.get(sessionId);
                        if (existingTimer) {
                            clearTimeout(existingTimer);
                            streamingCooldownTimers.delete(sessionId);
                        }

                        const timeoutId = setTimeout(() => {
                            set((state) => {
                                const memoryState = state.sessionMemoryState.get(sessionId);
                                if (!memoryState) return state;

                                if (memoryState.isStreaming) {
                                    return state;
                                }

                                const nextMemoryState = new Map(state.sessionMemoryState);
                                const { streamingCooldownUntil: _streamingCooldownUntil, ...rest } = memoryState;
                                void _streamingCooldownUntil;
                                nextMemoryState.set(sessionId, rest as SessionMemoryState);
                                return { sessionMemoryState: nextMemoryState };
                            });
                            streamingCooldownTimers.delete(sessionId);
                        }, 2000);

                        streamingCooldownTimers.set(sessionId, timeoutId);
                    }
                },

                syncMessages: (sessionId: string, messages: { info: Message; parts: Part[] }[]) => {
                    // Filter out reverted messages first
                    const revertMessageId = getSessionRevertMessageId(sessionId);
                    const messagesWithoutReverted = filterRevertedMessages(messages, revertMessageId);

                    const watermark = get().sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;
                    const messagesFiltered = watermark
                        ? messagesWithoutReverted.filter((message) => {
                              const messageId = message?.info?.id;
                              if (!messageId) return true;

                              return isIdNewer(messageId, watermark);
                          })
                        : messagesWithoutReverted;

                    set((state) => {
                        const newMessages = new Map(state.messages);
                        const previousMessages = state.messages.get(sessionId) || [];
                        const previousMessagesById = new Map(
                            previousMessages
                                .filter((msg) => typeof msg.info?.id === "string")
                                .map((msg) => [msg.info.id as string, msg])
                        );

                        const normalizedMessages = messagesFiltered.map((message) => {
                            const infoWithMarker = {
                                ...message.info,
                                clientRole: (message.info as any)?.clientRole ?? message.info.role,
                                userMessageMarker: message.info.role === "user" ? true : (message.info as any)?.userMessageMarker,
                                animationSettled:
                                    message.info.role === "assistant"
                                        ? (message.info as any)?.animationSettled ?? true
                                        : (message.info as any)?.animationSettled,
                            } as any;

                            const serverParts = (Array.isArray(message.parts) ? message.parts : []).map((part) => {
                                if (part?.type === 'text') {
                                    const raw = (part as any).text ?? (part as any).content ?? '';
                                    if (isExecutionForkMetaText(raw)) {
                                        return { ...part, synthetic: true } as Part;
                                    }
                                }
                                return part;
                            });
                            const messageId = typeof infoWithMarker?.id === "string" ? (infoWithMarker.id as string) : undefined;
                            const existingEntry = messageId ? previousMessagesById.get(messageId) : undefined;

                            if (existingEntry && existingEntry.info.role === "assistant") {
                                const existingParts = Array.isArray(existingEntry.parts) ? existingEntry.parts : [];
                                const existingLen = computePartsTextLength(existingParts);
                                const serverLen = computePartsTextLength(serverParts);
                                const storeHasStop = hasStopReasonStop(existingParts);

                                if (storeHasStop && existingLen > serverLen) {
                                    const mergedParts = mergePreferExistingParts(existingParts, serverParts);
                                    return {
                                        ...message,
                                        info: infoWithMarker,
                                        parts: mergedParts,
                                    };
                                }
                            }

                            return {
                                ...message,
                                info: infoWithMarker,
                                parts: serverParts,
                            };
                        });

                        const mergedMessages = dedupeMessagesById(normalizedMessages);

                        const previousIds = new Set(previousMessages.map((msg) => msg.info.id));
                        const nextIds = new Set(mergedMessages.map((msg) => msg.info.id));
                        const removedIds: string[] = [];
                        previousIds.forEach((id) => {
                            if (!nextIds.has(id)) {
                                removedIds.push(id);
                            }
                        });

                        newMessages.set(sessionId, mergedMessages);

                        const result: Record<string, any> = {
                            messages: newMessages,
                            isSyncing: true,
                        };

                        clearLifecycleTimersForIds(removedIds);
                        const updatedLifecycle = removeLifecycleEntries(state.messageStreamStates, removedIds);
                        if (updatedLifecycle !== state.messageStreamStates) {
                            result.messageStreamStates = updatedLifecycle;
                        }

                        if (removedIds.length > 0) {
                            const currentStreaming = state.streamingMessageIds.get(sessionId);
                            if (currentStreaming && removedIds.includes(currentStreaming)) {
                                result.streamingMessageIds = setStreamingIdForSession(
                                    result.streamingMessageIds ?? state.streamingMessageIds,
                                    sessionId,
                                    null
                                );
                            }
                        }

                        if (removedIds.length > 0) {
                            const nextIndex = removeMessageSessionIndexEntries(
                                result.messageSessionIndex ?? state.messageSessionIndex,
                                removedIds
                            );
                            if (nextIndex !== (result.messageSessionIndex ?? state.messageSessionIndex)) {
                                result.messageSessionIndex = nextIndex;
                            }
                        }

                        if (removedIds.length > 0) {
                            const nextPendingParts = new Map(state.pendingAssistantParts);
                            let pendingChanged = false;
                            removedIds.forEach((id) => {
                                if (nextPendingParts.delete(id)) {
                                    pendingChanged = true;
                                }
                            });
                            if (pendingChanged) {
                                result.pendingAssistantParts = nextPendingParts;
                            }
                        }

                        const targetIndex = result.messageSessionIndex ?? state.messageSessionIndex;
                        let indexAccumulator = targetIndex;
                        mergedMessages.forEach((message) => {
                            const id = (message?.info as { id?: unknown })?.id;
                            if (typeof id === "string" && id.length > 0) {
                                indexAccumulator = upsertMessageSessionIndex(indexAccumulator, id, sessionId);
                            }
                        });
                        if (indexAccumulator !== targetIndex) {
                            result.messageSessionIndex = indexAccumulator;
                        }

                        return result;
                    });

                    setTimeout(() => {
                        set({ isSyncing: false });
                    }, 100);
                },

                updateSessionCompaction: (sessionId: string, compactingTimestamp: number | null | undefined) => {
                    set((state) => {
                        const nextCompaction = new Map(state.sessionCompactionUntil);

                        if (!compactingTimestamp || compactingTimestamp <= 0) {
                            if (!nextCompaction.has(sessionId)) {
                                return state;
                            }
                            nextCompaction.delete(sessionId);
                            return { sessionCompactionUntil: nextCompaction };
                        }

                        const deadline = compactingTimestamp + COMPACTION_WINDOW_MS;
                        const existingDeadline = nextCompaction.get(sessionId);
                        if (existingDeadline === deadline) {
                            return state;
                        }

                        nextCompaction.set(sessionId, deadline);
                        return { sessionCompactionUntil: nextCompaction };
                    });
                },

                acknowledgeSessionAbort: (sessionId: string) => {
                    if (!sessionId) {
                        return;
                    }

                    set((state) => {
                        const record = state.sessionAbortFlags.get(sessionId);
                        if (!record || record.acknowledged) {
                            return state;
                        }

                        const nextAbortFlags = new Map(state.sessionAbortFlags);
                        nextAbortFlags.set(sessionId, { ...record, acknowledged: true });
                        return { sessionAbortFlags: nextAbortFlags } as Partial<MessageState>;
                    });
                },

                updateViewportAnchor: (sessionId: string, anchor: number) => {
                    set((state) => {
                        const memoryState = state.sessionMemoryState.get(sessionId) || {
                            viewportAnchor: 0,
                            isStreaming: false,
                            lastAccessedAt: Date.now(),
                            backgroundMessageCount: 0,
                        };

                        const newMemoryState = new Map(state.sessionMemoryState);
                        newMemoryState.set(sessionId, { ...memoryState, viewportAnchor: anchor });
                        return { sessionMemoryState: newMemoryState };
                    });
                },

                trimToViewportWindow: (sessionId: string, targetSize: number = MEMORY_LIMITS.VIEWPORT_MESSAGES, currentSessionId?: string) => {
                    const state = get();
                    const sessionMessages = state.messages.get(sessionId);
                    if (!sessionMessages || sessionMessages.length <= targetSize) {
                        return;
                    }

                    const memoryState = state.sessionMemoryState.get(sessionId) || {
                        viewportAnchor: sessionMessages.length - 1,
                        isStreaming: false,
                        lastAccessedAt: Date.now(),
                        backgroundMessageCount: 0,
                    };

                    if (memoryState.isStreaming && sessionId === currentSessionId) {
                        return;
                    }

                    const anchor = memoryState.viewportAnchor || sessionMessages.length - 1;
                    let start = Math.max(0, anchor - Math.floor(targetSize / 2));
                    const end = Math.min(sessionMessages.length, start + targetSize);

                    if (end === sessionMessages.length && end - start < targetSize) {
                        start = Math.max(0, end - targetSize);
                    }

                    const trimmedMessages = sessionMessages.slice(start, end);
                    const removedOlder = sessionMessages.slice(0, start);
                    const trimmedIds = new Set(trimmedMessages.map((message) => message.info.id));
                    const removedIds = sessionMessages
                        .filter((message) => !trimmedIds.has(message.info.id))
                        .map((message) => message.info.id);

                    set((state) => {
                        const newMessages = new Map(state.messages);
                        newMessages.set(sessionId, trimmedMessages);

                        const newMemoryState = new Map(state.sessionMemoryState);
                        const updatedMemoryState = {
                            ...memoryState,
                            viewportAnchor: anchor - start,
                            trimmedHeadMaxId: computeMaxTrimmedHeadId(removedOlder, memoryState.trimmedHeadMaxId),
                        };
                        newMemoryState.set(sessionId, updatedMemoryState);

                        const result: Record<string, any> = {
                            messages: newMessages,
                            sessionMemoryState: newMemoryState,
                        };

                        clearLifecycleTimersForIds(removedIds);
                        const updatedLifecycle = removeLifecycleEntries(state.messageStreamStates, removedIds);
                        if (updatedLifecycle !== state.messageStreamStates) {
                            result.messageStreamStates = updatedLifecycle;
                        }

                        if (removedIds.length > 0) {
                            const currentStreaming = state.streamingMessageIds.get(sessionId);
                            if (currentStreaming && removedIds.includes(currentStreaming)) {
                                result.streamingMessageIds = setStreamingIdForSession(
                                    result.streamingMessageIds ?? state.streamingMessageIds,
                                    sessionId,
                                    null
                                );
                            }
                        }

                        if (removedIds.length > 0) {
                            const nextIndex = removeMessageSessionIndexEntries(
                                result.messageSessionIndex ?? state.messageSessionIndex,
                                removedIds
                            );
                            if (nextIndex !== (result.messageSessionIndex ?? state.messageSessionIndex)) {
                                result.messageSessionIndex = nextIndex;
                            }
                        }

                        const targetIndex = result.messageSessionIndex ?? state.messageSessionIndex;
                        let indexAccumulator = targetIndex;
                        trimmedMessages.forEach((message) => {
                            const id = (message?.info as { id?: unknown })?.id;
                            if (typeof id === "string" && id.length > 0) {
                                indexAccumulator = upsertMessageSessionIndex(indexAccumulator, id, sessionId);
                            }
                        });
                        if (indexAccumulator !== targetIndex) {
                            result.messageSessionIndex = indexAccumulator;
                        }

                        return result;
                    });
                },

                evictLeastRecentlyUsed: (currentSessionId?: string) => {
                    const state = get();
                    const sessionCount = state.messages.size;

                    if (sessionCount <= MEMORY_LIMITS.MAX_SESSIONS) return;

                    const sessionsWithMemory: Array<[string, SessionMemoryState]> = [];
                    state.messages.forEach((_, sessionId) => {
                        const memoryState = state.sessionMemoryState.get(sessionId) || {
                            viewportAnchor: 0,
                            isStreaming: false,
                            lastAccessedAt: 0,
                            backgroundMessageCount: 0,
                        };
                        sessionsWithMemory.push([sessionId, memoryState]);
                    });

                    const evictable = sessionsWithMemory.filter(([id, memState]) => id !== currentSessionId && !memState.isStreaming);

                    if (evictable.length === 0) return;

                    evictable.sort((a, b) => a[1].lastAccessedAt - b[1].lastAccessedAt);
                    const lruSessionId = evictable[0][0];

                    set((state) => {
                        const removedMessages = state.messages.get(lruSessionId) || [];
                        const removedIds = removedMessages.map((message) => message.info.id);

                        const newMessages = new Map(state.messages);
                        const newMemoryState = new Map(state.sessionMemoryState);

                        newMessages.delete(lruSessionId);
                        newMemoryState.delete(lruSessionId);

                        const result: Record<string, any> = {
                            messages: newMessages,
                            sessionMemoryState: newMemoryState,
                        };

                        const nextPendingParts = new Map(state.pendingAssistantParts);
                        let pendingChanged = false;
                        nextPendingParts.forEach((entry, messageId) => {
                            if (entry.sessionId === lruSessionId) {
                                nextPendingParts.delete(messageId);
                                pendingChanged = true;
                            }
                        });
                        if (pendingChanged) {
                            result.pendingAssistantParts = nextPendingParts;
                        }

                        const nextCompaction = new Map(state.sessionCompactionUntil);
                        if (nextCompaction.delete(lruSessionId)) {
                            result.sessionCompactionUntil = nextCompaction;
                        }

                        clearLifecycleTimersForIds(removedIds);
                        const updatedLifecycle = removeLifecycleEntries(state.messageStreamStates, removedIds);
                        if (updatedLifecycle !== state.messageStreamStates) {
                            result.messageStreamStates = updatedLifecycle;
                        }

                        const nextIndex = removeMessageSessionIndexEntries(
                            result.messageSessionIndex ?? state.messageSessionIndex,
                            removedIds
                        );
                        if (nextIndex !== (result.messageSessionIndex ?? state.messageSessionIndex)) {
                            result.messageSessionIndex = nextIndex;
                        }

                        const nextStreamingIds = setStreamingIdForSession(state.streamingMessageIds, lruSessionId, null);
                        if (nextStreamingIds !== state.streamingMessageIds) {
                            result.streamingMessageIds = nextStreamingIds;
                        }

                        if (state.abortControllers.has(lruSessionId)) {
                            const nextControllers = new Map(state.abortControllers);
                            nextControllers.delete(lruSessionId);
                            result.abortControllers = nextControllers;
                        }

                        return result;
                    });
                },

                loadMoreMessages: async (sessionId: string, direction: "up" | "down" = "up") => {
                    const state = get();
                    const currentMessages = state.messages.get(sessionId);
                    const memoryState = state.sessionMemoryState.get(sessionId);

                    if (!currentMessages || !memoryState) {
                        return;
                    }

                    if (memoryState.totalAvailableMessages && currentMessages.length >= memoryState.totalAvailableMessages) {
                        return;
                    }

                        const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId));

                        if (direction === "up" && currentMessages.length > 0) {

                            const firstCurrentMessage = currentMessages[0];
                            const indexInAll = allMessages.findIndex((m) => m.info.id === firstCurrentMessage.info.id);

                            if (indexInAll > 0) {

                                const loadCount = Math.min(MEMORY_LIMITS.VIEWPORT_MESSAGES, indexInAll);
                                const newMessages = allMessages.slice(indexInAll - loadCount, indexInAll);

                                set((state) => {
                                    const updatedMessages = [...newMessages, ...currentMessages];
                                    const dedupedMessages = dedupeMessagesById(updatedMessages);
                                    const newMessagesMap = new Map(state.messages);
                                    newMessagesMap.set(sessionId, dedupedMessages);

                                    const addedCount = Math.max(0, dedupedMessages.length - currentMessages.length);
                                    const newMemoryState = new Map(state.sessionMemoryState);
                                    newMemoryState.set(sessionId, {
                                        ...memoryState,
                                        viewportAnchor: memoryState.viewportAnchor + addedCount,
                                        hasMoreAbove: indexInAll - loadCount > 0,
                                        totalAvailableMessages: allMessages.length,
                                    });

                                    return {
                                        messages: newMessagesMap,
                                        sessionMemoryState: newMemoryState,
                                    };
                                });
                            } else if (indexInAll === 0) {

                                set((state) => {
                                    const newMemoryState = new Map(state.sessionMemoryState);
                                    newMemoryState.set(sessionId, {
                                        ...memoryState,
                                        hasMoreAbove: false,
                                    });
                                    return { sessionMemoryState: newMemoryState };
                                });
                            }
                        }
                },

                getLastMessageModel: (sessionId: string) => {
                    const { messages } = get();
                    const sessionMessages = messages.get(sessionId);

                    if (!sessionMessages || sessionMessages.length === 0) {
                        return null;
                    }

                    for (let i = sessionMessages.length - 1; i >= 0; i--) {
                        const message = sessionMessages[i];
                        if (message.info.role === "assistant" && "providerID" in message.info && "modelID" in message.info) {
                            return {
                                providerID: (message.info as any).providerID,
                                modelID: (message.info as any).modelID,
                            };
                        }
                    }

                    return null;
                },
            }),
            {
                name: "message-store",
                storage: createJSONStorage(() => getSafeStorage()),
                partialize: (state: MessageStore) => ({
                    lastUsedProvider: state.lastUsedProvider,
                    sessionMemoryState: Array.from(state.sessionMemoryState.entries()).map(([sessionId, memory]) => [
                        sessionId,
                        {
                            viewportAnchor: memory.viewportAnchor,
                            isStreaming: memory.isStreaming,
                            lastAccessedAt: memory.lastAccessedAt,
                            backgroundMessageCount: memory.backgroundMessageCount,
                            totalAvailableMessages: memory.totalAvailableMessages,
                            hasMoreAbove: memory.hasMoreAbove,
                            trimmedHeadMaxId: memory.trimmedHeadMaxId,
                        },
                    ]),
                    sessionAbortFlags: Array.from(state.sessionAbortFlags.entries()).map(([sessionId, record]) => [
                        sessionId,
                        { timestamp: record.timestamp, acknowledged: record.acknowledged },
                    ]),
                }),
                merge: (persistedState: any, currentState: MessageStore): MessageStore => {
                    if (!persistedState) {
                        return currentState;
                    }

                    let restoredMemoryState = currentState.sessionMemoryState;
                    if (Array.isArray(persistedState.sessionMemoryState)) {
                        restoredMemoryState = new Map<string, SessionMemoryState>(
                            persistedState.sessionMemoryState.map((entry: [string, SessionMemoryState]) => entry)
                        );
                    }

                    let restoredAbortFlags = currentState.sessionAbortFlags;
                    if (Array.isArray(persistedState.sessionAbortFlags)) {
                        restoredAbortFlags = new Map<string, SessionAbortRecord>(persistedState.sessionAbortFlags);
                    }

                    return {
                        ...currentState,
                        lastUsedProvider: persistedState.lastUsedProvider ?? currentState.lastUsedProvider,
                        sessionMemoryState: restoredMemoryState,
                        sessionAbortFlags: restoredAbortFlags,
                    };
                },
            }
        ),
        {
            name: "message-store",
        }
    )
);