nextjs-chat-gpt/services/indexed-db/chat-store.ts

import type { UIMessage } from 'ai';

const DATABASE_NAME = 'chatgpt-agent-demo';
const DATABASE_VERSION = 1;
const CHATS_STORE = 'chats';
const MESSAGES_STORE = 'messages';

export interface StoredChat {
  id: string;
  title: string;
  created_at: string;
  updated_at: string;
}

interface StoredMessages {
  chatId: string;
  messages: UIMessage[];
  updated_at: string;
}

let databasePromise: Promise<IDBDatabase> | null = null;

function openDatabase(): Promise<IDBDatabase> {
  if (typeof window === 'undefined' || !('indexedDB' in window)) {
    return Promise.reject(new Error('IndexedDB is not available'));
  }

  if (databasePromise) return databasePromise;

  databasePromise = new Promise((resolve, reject) => {
    const request = indexedDB.open(DATABASE_NAME, DATABASE_VERSION);

    request.onupgradeneeded = () => {
      const database = request.result;

      if (!database.objectStoreNames.contains(CHATS_STORE)) {
        database.createObjectStore(CHATS_STORE, { keyPath: 'id' });
      }

      if (!database.objectStoreNames.contains(MESSAGES_STORE)) {
        database.createObjectStore(MESSAGES_STORE, { keyPath: 'chatId' });
      }
    };

    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
    request.onblocked = () =>
      reject(new Error('IndexedDB upgrade was blocked by another tab'));
  });

  return databasePromise;
}

function requestToPromise<T>(request: IDBRequest<T>): Promise<T> {
  return new Promise((resolve, reject) => {
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

function transactionToPromise(transaction: IDBTransaction): Promise<void> {
  return new Promise((resolve, reject) => {
    transaction.oncomplete = () => resolve();
    transaction.onerror = () => reject(transaction.error);
    transaction.onabort = () =>
      reject(transaction.error ?? new Error('IndexedDB transaction aborted'));
  });
}

export async function createChat(id: string, title: string): Promise<StoredChat> {
  const database = await openDatabase();
  const now = new Date().toISOString();
  const chat: StoredChat = {
    id,
    title,
    created_at: now,
    updated_at: now,
  };
  const transaction = database.transaction(CHATS_STORE, 'readwrite');
  transaction.objectStore(CHATS_STORE).put(chat);
  await transactionToPromise(transaction);
  return chat;
}

export async function listChats(): Promise<StoredChat[]> {
  const database = await openDatabase();
  const transaction = database.transaction(CHATS_STORE, 'readonly');
  const chats = await requestToPromise(
    transaction.objectStore(CHATS_STORE).getAll() as IDBRequest<StoredChat[]>
  );

  return chats.sort((a, b) => b.updated_at.localeCompare(a.updated_at));
}

export async function getMessages(chatId: string): Promise<UIMessage[]> {
  const database = await openDatabase();
  const transaction = database.transaction(MESSAGES_STORE, 'readonly');
  const record = await requestToPromise(
    transaction.objectStore(MESSAGES_STORE).get(chatId) as IDBRequest<
      StoredMessages | undefined
    >
  );

  return record?.messages ?? [];
}

export async function saveMessages(
  chatId: string,
  messages: UIMessage[]
): Promise<void> {
  const database = await openDatabase();
  const now = new Date().toISOString();
  const transaction = database.transaction(
    [CHATS_STORE, MESSAGES_STORE],
    'readwrite'
  );
  const chatsStore = transaction.objectStore(CHATS_STORE);

  const chatRequest = chatsStore.get(chatId) as IDBRequest<
    StoredChat | undefined
  >;
  chatRequest.onsuccess = () => {
    if (chatRequest.result) {
      chatsStore.put({ ...chatRequest.result, updated_at: now });
    }
  };

  transaction.objectStore(MESSAGES_STORE).put({
    chatId,
    messages,
    updated_at: now,
  } satisfies StoredMessages);

  await transactionToPromise(transaction);
}

export async function deleteChat(chatId: string): Promise<void> {
  const database = await openDatabase();
  const transaction = database.transaction(
    [CHATS_STORE, MESSAGES_STORE],
    'readwrite'
  );
  transaction.objectStore(CHATS_STORE).delete(chatId);
  transaction.objectStore(MESSAGES_STORE).delete(chatId);
  await transactionToPromise(transaction);
}