Add model versioning to storage and cache invalidation

src/contexts/MaiaEngineContext.tsx

@@ -34,6 +34,7 @@ export const MaiaEngineContextProvider: React.FC<{ children: ReactNode }> = ({
       model:
         process.env.NEXT_PUBLIC_MAIA_MODEL_URL ??
         'https://raw.githubusercontent.com/CSSLab/maia-platform-frontend/e23a50e/public/maia2/maia_rapid.onnx',
+      modelVersion: process.env.NEXT_PUBLIC_MAIA_MODEL_VERSION ?? '1',
       setStatus: setStatus,
       setProgress: setProgress,
       setError: setError,

src/lib/engine/maia.ts

@@ -6,6 +6,7 @@ import { MaiaModelStorage } from './storage'
 
 interface MaiaOptions {
   model: string
+  modelVersion: string
   setStatus: (status: MaiaStatus) => void
   setProgress: (progress: number) => void
   setError: (error: string) => void
@@ -14,11 +15,13 @@ interface MaiaOptions {
 class Maia {
   private model!: InferenceSession
   private modelUrl: string
+  private modelVersion: string
   private options: MaiaOptions
   private storage: MaiaModelStorage
 
   constructor(options: MaiaOptions) {
     this.modelUrl = options.model
+    this.modelVersion = options.modelVersion
     this.options = options
     this.storage = new MaiaModelStorage()
 
@@ -30,7 +33,7 @@ class Maia {
     await this.storage.requestPersistentStorage()
 
     console.log('Attempting to get model from IndexedDB...')
-    const buffer = await this.storage.getModel(this.modelUrl)
+    const buffer = await this.storage.getModel(this.modelUrl, this.modelVersion)
 
     if (buffer) {
       console.log('Model found in IndexedDB, initializing...')
@@ -94,7 +97,7 @@ class Maia {
       position += chunk.length
     }
 
-    await this.storage.storeModel(this.modelUrl, buffer.buffer)
+    await this.storage.storeModel(this.modelUrl, this.modelVersion, buffer.buffer)
 
     await this.initializeModel(buffer.buffer)
     this.options.setStatus('ready')

src/lib/engine/storage.ts

@@ -1,6 +1,7 @@
 interface ModelStorage {
   id: string
   url: string
+  version: string
   data: Blob
   timestamp: number
   size: number
@@ -34,7 +35,11 @@ export class MaiaModelStorage {
     })
   }
 
-  async storeModel(modelUrl: string, buffer: ArrayBuffer): Promise<void> {
+  async storeModel(
+    modelUrl: string,
+    modelVersion: string,
+    buffer: ArrayBuffer,
+  ): Promise<void> {
     try {
       const db = await this.openDB()
       const transaction = db.transaction([this.storeName], 'readwrite')
@@ -43,6 +48,7 @@ export class MaiaModelStorage {
       const modelData: ModelStorage = {
         id: 'maia-rapid-model',
         url: modelUrl,
+        version: modelVersion,
         data: new Blob([buffer]),
         timestamp: Date.now(),
         size: buffer.byteLength,
@@ -61,7 +67,10 @@ export class MaiaModelStorage {
     }
   }
 
-  async getModel(modelUrl: string): Promise<ArrayBuffer | null> {
+  async getModel(
+    modelUrl: string,
+    modelVersion: string,
+  ): Promise<ArrayBuffer | null> {
     console.log('Storage: getModel called with URL:', modelUrl)
 
     try {
@@ -95,18 +104,36 @@ export class MaiaModelStorage {
         return null
       }
 
-      // Check if URL matches (for cache invalidation)
-      if (modelData.url !== modelUrl) {
-        console.log('Storage: Model URL changed, clearing old cache')
-        console.log('Storage: Cached URL:', modelData.url)
-        console.log('Storage: Requested URL:', modelUrl)
+      // Version mismatch = genuinely new model, must re-download.
+      if (modelData.version && modelData.version !== modelVersion) {
+        console.log('Storage: Model version changed, clearing old cache')
+        console.log('Storage: Cached version:', modelData.version)
+        console.log('Storage: Required version:', modelVersion)
         await this.deleteModel()
-        console.log(
-          'Storage: Old cache cleared, model needs to be re-downloaded',
-        )
         return null
       }
 
+      // URL changed but same version (e.g. moved hosts) — keep the cache,
+      // just update the stored URL so future loads match immediately.
+      if (modelData.url !== modelUrl || modelData.version !== modelVersion) {
+        console.log('Storage: Updating stored URL/version')
+        try {
+          const rwTx = db.transaction([this.storeName], 'readwrite')
+          const rwStore = rwTx.objectStore(this.storeName)
+          await new Promise<void>((resolve, reject) => {
+            const req = rwStore.put({
+              ...modelData,
+              url: modelUrl,
+              version: modelVersion,
+            })
+            req.onsuccess = () => resolve()
+            req.onerror = () => reject(req.error)
+          })
+        } catch (e) {
+          console.warn('Storage: Failed to update stored URL/version:', e)
+        }
+      }
+
       console.log('Storage: Converting Blob to ArrayBuffer...')
       // Convert Blob back to ArrayBuffer
       const buffer = await modelData.data.arrayBuffer()