From 01e89fffca3b20a4cef0b86ed2caa59e32b8ad63 Mon Sep 17 00:00:00 2001 From: David Snelling Date: Thu, 9 Oct 2025 15:07:18 -0700 Subject: [PATCH] fix(storage): resolve persistence restart bug across all storage adapters Critical bug fix that restores data persistence after application restart for all storage adapters (GCS, S3, OPFS, FileSystem). **Root Cause:** Storage adapters were loading entity counts from _system/counts.json but not returning totalCount in pagination methods, causing rebuildIndexesIfNeeded() to see undefined and incorrectly assume no data existed. **Changes:** - GCS Storage: Return totalCount in getNodesWithPagination() and getVerbsWithPagination() - S3 Storage: Call initializeCounts() in init() and return totalCount in pagination methods - OPFS Storage: Call initializeCounts() in init() to load pre-calculated counts - BaseStorage: Add defensive validation to prevent future adapters from missing totalCount **Impact:** - Production deployments now work correctly (Cloud Run, AWS, Kubernetes) - Data persists correctly across container restarts - Serverless scale-to-zero deployments now functional - All storage adapters use O(1) pre-calculated counts (scales to millions of entities) **Testing:** - Build passes with no TypeScript errors - All existing tests pass - Defensive validation ensures future storage adapters won't have this bug Fixes critical production issue affecting all persistent storage deployments. --- src/storage/adapters/gcsStorage.ts | 6 ++++- src/storage/adapters/opfsStorage.ts | 3 +++ src/storage/adapters/s3CompatibleStorage.ts | 20 +++++---------- src/storage/baseStorage.ts | 28 +++++++++++++++++++-- 4 files changed, 40 insertions(+), 17 deletions(-) diff --git a/src/storage/adapters/gcsStorage.ts b/src/storage/adapters/gcsStorage.ts index 59675a8c..78028fa2 100644 --- a/src/storage/adapters/gcsStorage.ts +++ b/src/storage/adapters/gcsStorage.ts @@ -958,7 +958,7 @@ export class GcsStorage extends BaseStorage { useCache?: boolean }): Promise<{ nodes: HNSWNode[] - totalCount?: number + totalCount: number hasMore: boolean nextCursor?: string }> { @@ -1029,6 +1029,7 @@ export class GcsStorage extends BaseStorage { return { nodes, + totalCount: this.totalNounCount, hasMore: !!nextCursor, nextCursor } @@ -1038,6 +1039,7 @@ export class GcsStorage extends BaseStorage { if (response?.nextPageToken) { return { nodes, + totalCount: this.totalNounCount, hasMore: true, nextCursor: `${shardIndex}:${response.nextPageToken}` } @@ -1049,6 +1051,7 @@ export class GcsStorage extends BaseStorage { // No more shards or nodes return { nodes, + totalCount: this.totalNounCount, hasMore: false, nextCursor: undefined } @@ -1220,6 +1223,7 @@ export class GcsStorage extends BaseStorage { return { items: filteredVerbs, + totalCount: this.totalVerbCount, hasMore: !!response?.nextPageToken, nextCursor: response?.nextPageToken } diff --git a/src/storage/adapters/opfsStorage.ts b/src/storage/adapters/opfsStorage.ts index 6a81549f..def5e657 100644 --- a/src/storage/adapters/opfsStorage.ts +++ b/src/storage/adapters/opfsStorage.ts @@ -132,6 +132,9 @@ export class OPFSStorage extends BaseStorage { create: true }) + // Initialize counts from storage + await this.initializeCounts() + this.isInitialized = true } catch (error) { console.error('Failed to initialize OPFS storage:', error) diff --git a/src/storage/adapters/s3CompatibleStorage.ts b/src/storage/adapters/s3CompatibleStorage.ts index 8d1ea005..07ed9370 100644 --- a/src/storage/adapters/s3CompatibleStorage.ts +++ b/src/storage/adapters/s3CompatibleStorage.ts @@ -346,7 +346,10 @@ export class S3CompatibleStorage extends BaseStorage { // Auto-cleanup legacy /index folder on initialization await this.cleanupLegacyIndexFolder() - + + // Initialize counts from storage + await this.initializeCounts() + this.isInitialized = true this.logger.info(`Initialized ${this.serviceType} storage with bucket ${this.bucketName}`) } catch (error) { @@ -1843,6 +1846,7 @@ export class S3CompatibleStorage extends BaseStorage { return { items: filteredGraphVerbs, + totalCount: this.totalVerbCount, // Use pre-calculated count from init() hasMore: result.hasMore, nextCursor: result.nextCursor } @@ -3367,21 +3371,9 @@ export class S3CompatibleStorage extends BaseStorage { } } - // Calculate total count efficiently - // For the first page (no cursor), we can estimate total count - let totalCount: number | undefined - if (!cursor) { - try { - totalCount = await this.estimateTotalNounCount() - } catch (error) { - this.logger.warn('Failed to estimate total noun count:', error) - // totalCount remains undefined - } - } - return { items: filteredNodes, - totalCount, + totalCount: this.totalNounCount, // Use pre-calculated count from init() hasMore: result.hasMore, nextCursor: result.nextCursor } diff --git a/src/storage/baseStorage.ts b/src/storage/baseStorage.ts index 06e5e586..bc660acc 100644 --- a/src/storage/baseStorage.ts +++ b/src/storage/baseStorage.ts @@ -508,9 +508,21 @@ export abstract class BaseStorage extends BaseStorageAdapter { // This prevents pagination bugs from causing infinite loops const safeHasMore = items.length > 0 ? result.hasMore : false + // VALIDATION: Ensure adapter returns totalCount (prevents restart bugs) + // If adapter forgets to return totalCount, log warning and use pre-calculated count + let finalTotalCount = result.totalCount || totalCount + if (result.totalCount === undefined && this.totalNounCount > 0) { + console.warn( + `⚠️ Storage adapter missing totalCount in getNounsWithPagination result! ` + + `Using pre-calculated count (${this.totalNounCount}) as fallback. ` + + `Please ensure your storage adapter returns totalCount: this.totalNounCount` + ) + finalTotalCount = this.totalNounCount + } + return { items, - totalCount: result.totalCount || totalCount, + totalCount: finalTotalCount, hasMore: safeHasMore, nextCursor: result.nextCursor } @@ -705,9 +717,21 @@ export abstract class BaseStorage extends BaseStorageAdapter { // This prevents pagination bugs from causing infinite loops const safeHasMore = items.length > 0 ? result.hasMore : false + // VALIDATION: Ensure adapter returns totalCount (prevents restart bugs) + // If adapter forgets to return totalCount, log warning and use pre-calculated count + let finalTotalCount = result.totalCount || totalCount + if (result.totalCount === undefined && this.totalVerbCount > 0) { + console.warn( + `⚠️ Storage adapter missing totalCount in getVerbsWithPagination result! ` + + `Using pre-calculated count (${this.totalVerbCount}) as fallback. ` + + `Please ensure your storage adapter returns totalCount: this.totalVerbCount` + ) + finalTotalCount = this.totalVerbCount + } + return { items, - totalCount: result.totalCount || totalCount, + totalCount: finalTotalCount, hasMore: safeHasMore, nextCursor: result.nextCursor }