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 }