feat: server-side client state persistence

Author: psychedelicious

invokeai/app/api/dependencies.py

@@ -10,6 +10,7 @@
 from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage
 from invokeai.app.services.boards.boards_default import BoardService
 from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService
+from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite
 from invokeai.app.services.config.config_default import InvokeAIAppConfig
 from invokeai.app.services.download.download_default import DownloadQueueService
 from invokeai.app.services.events.events_fastapievents import FastAPIEventService
@@ -151,6 +152,7 @@ def initialize(
         style_preset_records = SqliteStylePresetRecordsStorage(db=db)
         style_preset_image_files = StylePresetImageFileStorageDisk(style_presets_folder / "images")
         workflow_thumbnails = WorkflowThumbnailFileStorageDisk(workflow_thumbnails_folder)
+        client_state_persistence = ClientStatePersistenceSqlite(db=db)
 
         services = InvocationServices(
             board_image_records=board_image_records,
@@ -181,6 +183,7 @@ def initialize(
             style_preset_records=style_preset_records,
             style_preset_image_files=style_preset_image_files,
             workflow_thumbnails=workflow_thumbnails,
+            client_state_persistence=client_state_persistence,
         )
 
         ApiDependencies.invoker = Invoker(services)

invokeai/app/api/routers/app_info.py

@@ -5,9 +5,9 @@
 from typing import Optional
 
 import torch
-from fastapi import Body
+from fastapi import Body, HTTPException, Query
 from fastapi.routing import APIRouter
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, JsonValue
 
 from invokeai.app.api.dependencies import ApiDependencies
 from invokeai.app.invocations.upscale import ESRGAN_MODELS
@@ -173,3 +173,50 @@ async def disable_invocation_cache() -> None:
 async def get_invocation_cache_status() -> InvocationCacheStatus:
     """Clears the invocation cache"""
     return ApiDependencies.invoker.services.invocation_cache.get_status()
+
+
+@app_router.get(
+    "/client_state",
+    operation_id="get_client_state_by_key",
+    response_model=JsonValue | None,
+)
+async def get_client_state_by_key(
+    key: str = Query(..., description="Key to retrieve from client state persistence"),
+) -> JsonValue | None:
+    """Gets the client state"""
+    try:
+        return ApiDependencies.invoker.services.client_state_persistence.get_by_key(key)
+    except Exception as e:
+        logging.error(f"Error getting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error setting client state")
+
+
+@app_router.post(
+    "/client_state",
+    operation_id="set_client_state",
+    response_model=None,
+)
+async def set_client_state(
+    key: str = Body(..., description="Key to set"),
+    value: JsonValue = Body(..., description="Value of the key"),
+) -> None:
+    """Sets the client state"""
+    try:
+        ApiDependencies.invoker.services.client_state_persistence.set_by_key(key, value)
+    except Exception as e:
+        logging.error(f"Error setting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error setting client state")
+
+
+@app_router.delete(
+    "/client_state",
+    operation_id="delete_client_state",
+    responses={204: {"description": "Client state deleted"}},
+)
+async def delete_client_state() -> None:
+    """Deletes the client state"""
+    try:
+        ApiDependencies.invoker.services.client_state_persistence.delete()
+    except Exception as e:
+        logging.error(f"Error deleting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error deleting client state")

invokeai/app/services/client_state_persistence/client_state_persistence_base.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+
+from pydantic import JsonValue
+
+
+class ClientStatePersistenceABC(ABC):
+    """
+    Base class for client persistence implementations.
+    This class defines the interface for persisting client data.
+    """
+
+    @abstractmethod
+    def set_by_key(self, key: str, value: JsonValue) -> None:
+        """
+        Store the data for the client.
+
+        :param data: The client data to be stored.
+        """
+        pass
+
+    @abstractmethod
+    def get_by_key(self, key: str) -> JsonValue | None:
+        """
+        Get the data for the client.
+
+        :return: The client data.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self) -> None:
+        """
+        Delete the data for the client.
+        """
+        pass

invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py

@@ -0,0 +1,65 @@
+import json
+
+from pydantic import JsonValue
+
+from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
+from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+
+
+class ClientStatePersistenceSqlite(ClientStatePersistenceABC):
+    """
+    Base class for client persistence implementations.
+    This class defines the interface for persisting client data.
+    """
+
+    def __init__(self, db: SqliteDatabase) -> None:
+        super().__init__()
+        self._db = db
+        self._default_row_id = 1
+
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+
+    def set_by_key(self, key: str, value: JsonValue) -> None:
+        state = self.get() or {}
+        state.update({key: value})
+
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                INSERT INTO client_state (id, data)
+                VALUES ({self._default_row_id}, ?)
+                ON CONFLICT(id) DO UPDATE
+                  SET data = excluded.data;
+                """,
+                (json.dumps(state),),
+            )
+
+    def get(self) -> dict[str, JsonValue] | None:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                SELECT data FROM client_state
+                WHERE id = {self._default_row_id}
+                """
+            )
+            row = cursor.fetchone()
+            if row is None:
+                return None
+            return json.loads(row[0])
+
+    def get_by_key(self, key: str) -> JsonValue | None:
+        state = self.get()
+        if state is None:
+            return None
+        return state.get(key, None)
+
+    def delete(self) -> None:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                DELETE FROM client_state
+                WHERE id = {self._default_row_id}
+                """
+            )

invokeai/app/services/invocation_services.py

@@ -17,6 +17,7 @@
     from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
     from invokeai.app.services.boards.boards_base import BoardServiceABC
     from invokeai.app.services.bulk_download.bulk_download_base import BulkDownloadBase
+    from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
     from invokeai.app.services.config import InvokeAIAppConfig
     from invokeai.app.services.download import DownloadQueueServiceBase
     from invokeai.app.services.events.events_base import EventServiceBase
@@ -73,6 +74,7 @@ def __init__(
         style_preset_records: "StylePresetRecordsStorageBase",
         style_preset_image_files: "StylePresetImageFileStorageBase",
         workflow_thumbnails: "WorkflowThumbnailServiceBase",
+        client_state_persistence: "ClientStatePersistenceABC",
     ):
         self.board_images = board_images
         self.board_image_records = board_image_records
@@ -102,3 +104,4 @@ def __init__(
         self.style_preset_records = style_preset_records
         self.style_preset_image_files = style_preset_image_files
         self.workflow_thumbnails = workflow_thumbnails
+        self.client_state_persistence = client_state_persistence

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -23,6 +23,7 @@
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_18 import build_migration_18
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_19 import build_migration_19
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_20 import build_migration_20
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import build_migration_21
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -63,6 +64,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_18())
     migrator.register_migration(build_migration_19(app_config=config))
     migrator.register_migration(build_migration_20())
+    migrator.register_migration(build_migration_21())
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_21.py

@@ -0,0 +1,40 @@
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration21Callback:
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute(
+            """
+            CREATE TABLE client_state (
+              id          INTEGER PRIMARY KEY CHECK(id = 1),
+              data        TEXT    NOT NULL, -- Frontend will handle the shape of this data
+              updated_at  DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP)
+            );
+            """
+        )
+        cursor.execute(
+            """
+            CREATE TRIGGER tg_client_state_updated_at
+            AFTER UPDATE ON client_state
+            FOR EACH ROW
+            BEGIN
+              UPDATE client_state
+                SET updated_at = CURRENT_TIMESTAMP
+              WHERE id = OLD.id;
+            END;
+            """
+        )
+
+
+def build_migration_21() -> Migration:
+    """Builds the migration object for migrating from version 20 to version 21. This includes:
+    - Creating the `client_state` table.
+    - Adding a trigger to update the `updated_at` field on updates.
+    """
+    return Migration(
+        from_version=20,
+        to_version=21,
+        callback=Migration21Callback(),
+    )

invokeai/frontend/web/src/app/store/enhancers/reduxRemember/driver.ts

@@ -1,9 +1,12 @@
 import { StorageError } from 'app/store/enhancers/reduxRemember/errors';
+import { $authToken } from 'app/store/nanostores/authToken';
 import { $projectId } from 'app/store/nanostores/projectId';
 import type { UseStore } from 'idb-keyval';
 import { clear, createStore as createIDBKeyValStore, get, set } from 'idb-keyval';
 import { atom } from 'nanostores';
 import type { Driver } from 'redux-remember';
+import { getBaseUrl } from 'services/api';
+import { buildAppInfoUrl } from 'services/api/endpoints/appInfo';
 
 // Create a custom idb-keyval store (just needed to customize the name)
 const $idbKeyValStore = atom<UseStore>(createIDBKeyValStore('invoke', 'invoke-store'));
@@ -38,3 +41,73 @@ export const idbKeyValDriver: Driver = {
     }
   },
 };
+
+const getHeaders = (extra?: Record<string, string>) => {
+  const headers = new Headers();
+  const authToken = $authToken.get();
+  if (authToken) {
+    headers.set('Authorization', `Bearer ${authToken}`);
+  }
+  const projectId = $projectId.get();
+  if (projectId) {
+    headers.set('project-id', projectId);
+  }
+  for (const [key, value] of Object.entries(extra ?? {})) {
+    headers.set(key, value);
+  }
+  return headers;
+};
+
+export const serverBackedDriver: Driver = {
+  getItem: async (key) => {
+    try {
+      const baseUrl = getBaseUrl();
+      const path = buildAppInfoUrl('client_state', { key });
+      const url = `${baseUrl}/${path}`;
+      const headers = getHeaders();
+      const res = await fetch(url, { headers, method: 'GET' });
+      if (!res.ok) {
+        throw new Error(`Response status: ${res.status}`);
+      }
+      const json = await res.json();
+      return json;
+    } catch (originalError) {
+      throw new StorageError({
+        key,
+        projectId: $projectId.get(),
+        originalError,
+      });
+    }
+  },
+  setItem: async (key, value) => {
+    try {
+      const baseUrl = getBaseUrl();
+      const path = buildAppInfoUrl('client_state');
+      const url = `${baseUrl}/${path}`;
+      const headers = getHeaders({ 'content-type': 'application/json' });
+      const res = await fetch(url, { headers, method: 'POST', body: JSON.stringify({ key, value }) });
+      if (!res.ok) {
+        throw new Error(`Response status: ${res.status}`);
+      }
+      return value;
+    } catch (originalError) {
+      throw new StorageError({
+        key,
+        value,
+        projectId: $projectId.get(),
+        originalError,
+      });
+    }
+  },
+};
+
+export const resetClientState = async () => {
+  const baseUrl = getBaseUrl();
+  const path = buildAppInfoUrl('client_state');
+  const url = `${baseUrl}/${path}`;
+  const headers = getHeaders();
+  const res = await fetch(url, { headers, method: 'DELETE' });
+  if (!res.ok) {
+    throw new Error(`Response status: ${res.status}`);
+  }
+};