From f59e203eb2b87335880f7a0751c5c7bc2f2ba213 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 12:35:42 -0700 Subject: [PATCH 01/10] feat(kvstore/ocdbt): add cache invalidation for manifest/btree/version Adds invalidateOcdbtCaches() so consumers can clear the three metadata caches after a server-side OCDBT mutation, forcing the next read to resolve a fresh root. Also removes the per-instance root memoization on OcdbtKvStore -- it was redundant with the ocdbt:version SimpleAsyncCache and prevented invalidation from taking effect. --- src/chunk_manager/generic_file_source.ts | 9 +++ src/kvstore/ocdbt/backend.ts | 27 +++----- src/kvstore/ocdbt/metadata_cache.ts | 85 +++++++++++++++++++++++- 3 files changed, 103 insertions(+), 18 deletions(-) diff --git a/src/chunk_manager/generic_file_source.ts b/src/chunk_manager/generic_file_source.ts index 45937f4a27..07cd99707d 100644 --- a/src/chunk_manager/generic_file_source.ts +++ b/src/chunk_manager/generic_file_source.ts @@ -69,6 +69,15 @@ export class SimpleAsyncCache extends ChunkSourceBase { progressOptions: ProgressOptions, ) => Promise<{ size: number; data: Value }>; + invalidate(key: Key) { + const encodedKey = this.encodeKeyFunction(key); + const chunk = this.chunks.get(encodedKey); + if (chunk !== undefined) { + chunk.freeSystemMemory(); + this.chunkManager.queueManager.updateChunkState(chunk, ChunkState.QUEUED); + } + } + get(key: Key, options: Partial): Promise { const encodedKey = this.encodeKeyFunction(key); let chunk = this.chunks.get(encodedKey); diff --git a/src/kvstore/ocdbt/backend.ts b/src/kvstore/ocdbt/backend.ts index 24486f57f8..371a244707 100644 --- a/src/kvstore/ocdbt/backend.ts +++ b/src/kvstore/ocdbt/backend.ts @@ -34,7 +34,6 @@ import { import { getRoot } from "#src/kvstore/ocdbt/read_version.js"; import { getOcdbtUrl } from "#src/kvstore/ocdbt/url.js"; import { type VersionSpecifier } from "#src/kvstore/ocdbt/version_specifier.js"; -import type { BtreeGenerationReference } from "#src/kvstore/ocdbt/version_tree.js"; import type { ProgressOptions } from "#src/util/progress_listener.js"; export class OcdbtKvStore implements KvStore { @@ -44,19 +43,13 @@ export class OcdbtKvStore implements KvStore { public version: VersionSpecifier | undefined, ) {} - private root: BtreeGenerationReference | undefined; - - private async getRoot(options: Partial) { - let { root } = this; - if (root === undefined) { - root = this.root = await getRoot( - this.sharedKvStoreContext, - this.baseUrl, - this.version, - options, - ); - } - return root; + private resolveRoot(options: Partial) { + return getRoot( + this.sharedKvStoreContext, + this.baseUrl, + this.version, + options, + ); } getUrl(key: string) { @@ -67,7 +60,7 @@ export class OcdbtKvStore implements KvStore { key: string, options: StatOptions, ): Promise { - const root = await this.getRoot(options); + const root = await this.resolveRoot(options); const encodedKey = new TextEncoder().encode(key) as Key; const entry = await findEntryInRoot( this.sharedKvStoreContext, @@ -85,7 +78,7 @@ export class OcdbtKvStore implements KvStore { key: string, options: DriverReadOptions, ): Promise { - const root = await this.getRoot(options); + const root = await this.resolveRoot(options); const encodedKey = new TextEncoder().encode(key) as Key; const entry = await findEntryInRoot( this.sharedKvStoreContext, @@ -105,7 +98,7 @@ export class OcdbtKvStore implements KvStore { prefix: string, options: DriverListOptions, ): Promise { - const root = await this.getRoot(options); + const root = await this.resolveRoot(options); const encodedPrefix = new TextEncoder().encode(prefix) as Key; return await listRoot( this.sharedKvStoreContext, diff --git a/src/kvstore/ocdbt/metadata_cache.ts b/src/kvstore/ocdbt/metadata_cache.ts index 361c048630..2c9e30236d 100644 --- a/src/kvstore/ocdbt/metadata_cache.ts +++ b/src/kvstore/ocdbt/metadata_cache.ts @@ -14,6 +14,7 @@ * limitations under the License. */ +import { ChunkState } from "#src/chunk_manager/base.js"; import { SimpleAsyncCache } from "#src/chunk_manager/generic_file_source.js"; import type { SharedKvStoreContextCounterpart } from "#src/kvstore/backend.js"; import type { BtreeNode } from "#src/kvstore/ocdbt/btree.js"; @@ -27,7 +28,12 @@ import type { ManifestWithVersionTree, } from "#src/kvstore/ocdbt/manifest.js"; import { decodeManifest } from "#src/kvstore/ocdbt/manifest.js"; -import type { VersionTreeNode } from "#src/kvstore/ocdbt/version_tree.js"; +import type { VersionSpecifier } from "#src/kvstore/ocdbt/version_specifier.js"; +import { formatVersion } from "#src/kvstore/ocdbt/version_specifier.js"; +import type { + BtreeGenerationReference, + VersionTreeNode, +} from "#src/kvstore/ocdbt/version_tree.js"; import { decodeVersionTreeNode } from "#src/kvstore/ocdbt/version_tree.js"; import { pipelineUrlJoin } from "#src/kvstore/url.js"; import type { ProgressOptions } from "#src/util/progress_listener.js"; @@ -84,6 +90,83 @@ export function getManifest( return cache.get(dataFile, options); } +export function invalidateOcdbtCaches( + sharedKvStoreContext: SharedKvStoreContextCounterpart, + _baseUrl: string, +) { + // Invalidate the cached manifest for this OCDBT database + const manifestCache = sharedKvStoreContext.chunkManager.memoize.get( + "ocdbt:manifest", + () => { + const cache = new SimpleAsyncCache( + sharedKvStoreContext.chunkManager.addRef(), + { + get: async () => { + throw new Error("unreachable"); + }, + }, + ); + cache.registerDisposer(sharedKvStoreContext.addRef()); + return cache; + }, + ); + for (const chunk of manifestCache.chunks.values()) { + chunk.freeSystemMemory(); + manifestCache.chunkManager.queueManager.updateChunkState( + chunk, + ChunkState.QUEUED, + ); + } + // Also invalidate all btree nodes since they may reference stale data + // after server-side mutations. This is broader than strictly necessary + // but btree nodes are small and fast to re-fetch. + const btreeCache = sharedKvStoreContext.chunkManager.memoize.get( + "ocdbt:btree", + () => + makeIndirectDataReferenceCache( + sharedKvStoreContext, + "b+tree node", + decodeBtreeNode, + ), + ); + for (const chunk of btreeCache.chunks.values()) { + chunk.freeSystemMemory(); + btreeCache.chunkManager.queueManager.updateChunkState( + chunk, + ChunkState.QUEUED, + ); + } + // Invalidate the cached BtreeGenerationReference so the next read + // resolves a fresh root from the updated manifest. + const versionCache = sharedKvStoreContext.chunkManager.memoize.get( + "ocdbt:version", + () => { + const cache = new SimpleAsyncCache< + { url: string; version: VersionSpecifier | undefined }, + BtreeGenerationReference + >(sharedKvStoreContext.chunkManager.addRef(), { + get: async () => { + throw new Error("unreachable"); + }, + encodeKey: ({ url, version }) => + JSON.stringify([ + url, + version !== undefined ? formatVersion(version) : undefined, + ]), + }); + cache.registerDisposer(sharedKvStoreContext.addRef()); + return cache; + }, + ); + for (const chunk of versionCache.chunks.values()) { + chunk.freeSystemMemory(); + versionCache.chunkManager.queueManager.updateChunkState( + chunk, + ChunkState.QUEUED, + ); + } +} + export async function getResolvedManifest( sharedKvStoreContext: SharedKvStoreContextCounterpart, url: string, From 99c82a79efc169b822bb0cdd844e8661f5eca475 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 12:45:44 -0700 Subject: [PATCH 02/10] feat(datasource/graphene): OCDBT segmentation support Adds ocdbt_seg / ocdbt_path graph-info fields. When ocdbt_seg is set, segmentation volume reads route through a per-scale OCDBT pipeline URL (scales auto-discovered via list()). Non-OCDBT scales are filtered from getSources() so the graphene layer only shows data available in the fork. After a multicut, invalidates the OCDBT metadata caches via RPC so split supervoxels become visible without a manual reload. Also skips the "supervoxel already selected" guard in the multicut tool when ocdbt_seg is active, since SV splits require selecting the same supervoxel on both sides of the cut. --- src/datasource/graphene/backend.ts | 7 ++ src/datasource/graphene/base.ts | 1 + src/datasource/graphene/frontend.ts | 174 ++++++++++++++++++++++++++-- 3 files changed, 173 insertions(+), 9 deletions(-) diff --git a/src/datasource/graphene/backend.ts b/src/datasource/graphene/backend.ts index 564aa3a76a..73997a80d4 100644 --- a/src/datasource/graphene/backend.ts +++ b/src/datasource/graphene/backend.ts @@ -28,6 +28,7 @@ import type { } from "#src/datasource/graphene/base.js"; import { getGrapheneFragmentKey, + GRAPHENE_INVALIDATE_OCDBT_RPC_ID, GRAPHENE_MESH_NEW_SEGMENT_RPC_ID, ChunkedGraphSourceParameters, MeshSourceParameters, @@ -42,6 +43,7 @@ import { decodeManifestChunk } from "#src/datasource/precomputed/backend.js"; import { WithSharedKvStoreContextCounterpart } from "#src/kvstore/backend.js"; import type { KvStoreWithPath, ReadResponse } from "#src/kvstore/index.js"; import { readKvStore } from "#src/kvstore/index.js"; +import { invalidateOcdbtCaches } from "#src/kvstore/ocdbt/metadata_cache.js"; import type { FragmentChunk, ManifestChunk } from "#src/mesh/backend.js"; import { assignMeshFragmentData, MeshSource } from "#src/mesh/backend.js"; import { decodeDraco } from "#src/mesh/draco/index.js"; @@ -643,3 +645,8 @@ registerRPC(GRAPHENE_MESH_NEW_SEGMENT_RPC_ID, function (x) { const obj = this.get(x.rpcId); obj.addNewSegment(x.segment); }); + +registerRPC(GRAPHENE_INVALIDATE_OCDBT_RPC_ID, function (x) { + const source = this.get(x.layerId) as GrapheneChunkedGraphChunkSource; + invalidateOcdbtCaches(source.sharedKvStoreContext, x.baseUrl); +}); diff --git a/src/datasource/graphene/base.ts b/src/datasource/graphene/base.ts index 844cddf352..7759660c29 100644 --- a/src/datasource/graphene/base.ts +++ b/src/datasource/graphene/base.ts @@ -32,6 +32,7 @@ import type { FetchOk, HttpError } from "#src/util/http_request.js"; export const PYCG_APP_VERSION = 1; export const GRAPHENE_MESH_NEW_SEGMENT_RPC_ID = "GrapheneMeshSource:NewSegment"; +export const GRAPHENE_INVALIDATE_OCDBT_RPC_ID = "Graphene:InvalidateOcdbt"; export enum VolumeChunkEncoding { RAW = 0, diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index a401e2efc3..b823227e7f 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -49,6 +49,7 @@ import { CHUNKED_GRAPH_RENDER_LAYER_UPDATE_SOURCES_RPC_ID, ChunkedGraphSourceParameters, getGrapheneFragmentKey, + GRAPHENE_INVALIDATE_OCDBT_RPC_ID, GRAPHENE_MESH_NEW_SEGMENT_RPC_ID, isBaseSegmentId, makeChunkedGraphChunkSpecification, @@ -73,6 +74,7 @@ import { getSegmentPropertyMap, parseMultiscaleVolumeInfo, PrecomputedMultiscaleVolumeChunkSource, + PrecomputedVolumeChunkSource, } from "#src/datasource/precomputed/frontend.js"; import { WithSharedKvStoreContext } from "#src/kvstore/chunk_source_frontend.js"; import type { SharedKvStoreContext } from "#src/kvstore/frontend.js"; @@ -132,6 +134,9 @@ import { } from "#src/sliceview/frontend.js"; import type { SliceViewRenderLayer } from "#src/sliceview/renderlayer.js"; import { SliceViewPanelRenderLayer } from "#src/sliceview/renderlayer.js"; +import type { VolumeSourceOptions } from "#src/sliceview/volume/base.js"; +import { makeDefaultVolumeChunkSpecifications } from "#src/sliceview/volume/base.js"; +import type { VolumeChunkSource } from "#src/sliceview/volume/frontend.js"; import { StatusMessage } from "#src/status.js"; import { TrackableBoolean, @@ -164,6 +169,7 @@ import { registerTool, } from "#src/ui/tool.js"; import { Uint64Set } from "#src/uint64_set.js"; +import { transposeNestedArrays } from "#src/util/array.js"; import { packColor } from "#src/util/color.js"; import type { Owned } from "#src/util/disposable.js"; import { RefCounted } from "#src/util/disposable.js"; @@ -281,6 +287,8 @@ const N_BITS_FOR_LAYER_ID_DEFAULT = 8; class GraphInfo { chunkSize: vec3; nBitsForLayerId: number; + ocdbtSeg: boolean; + ocdbtPath: string | undefined; constructor(obj: any) { verifyObject(obj); this.chunkSize = verifyObjectProperty(obj, "chunk_size", (x) => @@ -292,11 +300,25 @@ class GraphInfo { verifyPositiveInt, N_BITS_FOR_LAYER_ID_DEFAULT, ); + this.ocdbtSeg = verifyOptionalObjectProperty( + obj, + "ocdbt_seg", + verifyBoolean, + false, + ); + this.ocdbtPath = verifyOptionalObjectProperty( + obj, + "ocdbt_path", + verifyOptionalString, + undefined, + ); } } interface GrapheneMultiscaleVolumeInfo extends MultiscaleVolumeInfo { dataUrl: string; + ocdbtDataUrl: string | undefined; + ocdbtScales: Set; app: AppInfo; graph: GraphInfo; } @@ -309,15 +331,28 @@ function parseGrapheneMultiscaleVolumeInfo( const dataUrl = verifyObjectProperty(obj, "data_dir", verifyString); const app = verifyObjectProperty(obj, "app", (x) => new AppInfo(url, x)); const graph = verifyObjectProperty(obj, "graph", (x) => new GraphInfo(x)); + let ocdbtDataUrl: string | undefined; + if (graph.ocdbtSeg && graph.ocdbtPath) { + let ocdbtBase = dataUrl; + if (!ocdbtBase.endsWith("/")) ocdbtBase += "/"; + ocdbtBase += graph.ocdbtPath; + if (!ocdbtBase.endsWith("/")) ocdbtBase += "/"; + ocdbtDataUrl = `${ocdbtBase}|ocdbt:`; + } return { ...volumeInfo, app, graph, dataUrl, + ocdbtDataUrl, + ocdbtScales: new Set(), }; } class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChunkSource { + private volumeChunkSources: { invalidateCache(): void }[] = []; + private chunkedGraphChunkSource: GrapheneChunkedGraphChunkSource | undefined; + constructor( sharedKvStoreContext: SharedKvStoreContext, public info: GrapheneMultiscaleVolumeInfo, @@ -325,6 +360,107 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu super(sharedKvStoreContext, info.dataUrl, info); } + resolveScaleUrl(scaleKey: string): string { + const { ocdbtDataUrl, ocdbtScales } = this.info; + const baseUrl = + ocdbtDataUrl && ocdbtScales.has(scaleKey) ? ocdbtDataUrl : this.url; + return kvstoreEnsureDirectoryPipelineUrl( + this.sharedKvStoreContext.kvStoreContext.resolveRelativePath( + baseUrl, + scaleKey, + ), + ); + } + + getSources(volumeSourceOptions: VolumeSourceOptions) { + const modelResolution = this.info.scales[0].resolution; + const { rank } = this; + const sources = transposeNestedArrays( + this.info.scales + .filter((x) => !x.hidden) + .filter((x) => x.key !== "placeholder") + .filter( + (x) => + !this.info.graph.ocdbtSeg || this.info.ocdbtScales.has(x.key), + ) + .map((scaleInfo) => { + const { resolution } = scaleInfo; + const stride = rank + 1; + const chunkToMultiscaleTransform = new Float32Array(stride * stride); + chunkToMultiscaleTransform[chunkToMultiscaleTransform.length - 1] = 1; + const { lowerBounds: baseLowerBound, upperBounds: baseUpperBound } = + this.info.modelSpace.boundingBoxes[0].box; + const lowerClipBound = new Float32Array(rank); + const upperClipBound = new Float32Array(rank); + for (let i = 0; i < 3; ++i) { + const relativeScale = resolution[i] / modelResolution[i]; + chunkToMultiscaleTransform[stride * i + i] = relativeScale; + const voxelOffsetValue = scaleInfo.voxelOffset[i]; + chunkToMultiscaleTransform[stride * rank + i] = + voxelOffsetValue * relativeScale; + lowerClipBound[i] = + baseLowerBound[i] / relativeScale - voxelOffsetValue; + upperClipBound[i] = + baseUpperBound[i] / relativeScale - voxelOffsetValue; + } + if (rank === 4) { + chunkToMultiscaleTransform[stride * 3 + 3] = 1; + lowerClipBound[3] = baseLowerBound[3]; + upperClipBound[3] = baseUpperBound[3]; + } + return makeDefaultVolumeChunkSpecifications({ + rank, + dataType: this.dataType, + chunkToMultiscaleTransform, + upperVoxelBound: scaleInfo.size, + volumeType: this.volumeType, + chunkDataSizes: scaleInfo.chunkSizes, + baseVoxelOffset: scaleInfo.voxelOffset, + compressedSegmentationBlockSize: + scaleInfo.compressedSegmentationBlockSize, + volumeSourceOptions, + }).map( + (spec): SliceViewSingleResolutionSource => ({ + chunkSource: this.chunkManager.getChunkSource( + PrecomputedVolumeChunkSource, + { + sharedKvStoreContext: this.sharedKvStoreContext, + spec, + parameters: { + url: this.resolveScaleUrl(scaleInfo.key), + encoding: scaleInfo.encoding, + sharding: scaleInfo.sharding, + }, + }, + ), + chunkToMultiscaleTransform, + lowerClipBound, + upperClipBound, + }), + ); + }), + ); + this.volumeChunkSources = sources.flat().map((s) => s.chunkSource); + return sources; + } + + invalidateVolumeSources() { + // Invalidate OCDBT metadata caches first so that when volume chunks + // are re-queued and start downloading, they read fresh metadata. + if (this.info.graph.ocdbtSeg && this.chunkedGraphChunkSource?.rpc) { + this.chunkedGraphChunkSource.rpc.invoke( + GRAPHENE_INVALIDATE_OCDBT_RPC_ID, + { + layerId: this.chunkedGraphChunkSource.rpcId, + baseUrl: this.info.ocdbtDataUrl, + }, + ); + } + for (const source of this.volumeChunkSources) { + source.invalidateCache(); + } + } + getChunkedGraphSource() { const { rank } = this; const scaleInfo = this.info.scales[0]; @@ -352,15 +488,17 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu lowerClipBound[i] = baseLowerBound[i]; upperClipBound[i] = baseUpperBound[i]; } + const chunkSource = this.chunkManager.getChunkSource( + GrapheneChunkedGraphChunkSource, + { + spec, + sharedKvStoreContext: this.sharedKvStoreContext, + parameters: { url: `${this.info.app!.segmentationUrl}/node` }, + }, + ); + this.chunkedGraphChunkSource = chunkSource; return { - chunkSource: this.chunkManager.getChunkSource( - GrapheneChunkedGraphChunkSource, - { - spec, - sharedKvStoreContext: this.sharedKvStoreContext, - parameters: { url: `${this.info.app!.segmentationUrl}/node` }, - }, - ), + chunkSource, chunkToMultiscaleTransform, lowerClipBound, upperClipBound, @@ -613,6 +751,18 @@ async function getVolumeDataSource( stateJson: any, ): Promise { const info = parseGrapheneMultiscaleVolumeInfo(metadata, url); + if (info.ocdbtDataUrl) { + const listResult = await sharedKvStoreContext.kvStoreContext.list( + info.ocdbtDataUrl, + { responseKeys: "suffix", ...options }, + ); + const knownScaleKeys = new Set(info.scales.map((s) => s.key)); + for (const dir of listResult.directories) { + if (knownScaleKeys.has(dir)) { + info.ocdbtScales.add(dir); + } + } + } const volume = new GrapheneMultiscaleVolumeChunkSource( sharedKvStoreContext, info, @@ -1742,6 +1892,9 @@ class GraphConnection extends SegmentationGraphSourceConnection { const newValues = new Uint64Set(); newValues.add(splitRoots); this.state.replaceSegments(oldValues, newValues); + if (this.graph.info.graph.ocdbtSeg) { + this.chunkSource.invalidateVolumeSources(); + } return true; } } @@ -2891,7 +3044,10 @@ class MulticutSegmentsTool extends LayerTool { return; } const isRoot = rootId === segmentId; - if (!isRoot) { + // Supervoxel splits require selecting the same supervoxel on both + // sides of the cut (the split happens within one supervoxel), so + // skip the duplicate-selection guard when ocdbtSeg is active. + if (!isRoot && !graphConnection.graph.info.graph.ocdbtSeg) { for (const segment of segments) { if (segment === segmentId) { StatusMessage.showTemporaryMessage( From 6dca21f60d10d5398f6aae37ca27a836acee3658 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 14:04:54 -0700 Subject: [PATCH 03/10] feat(kvstore): add kvstack driver for key range-routed overlays Adds a base kvstore driver that routes reads/stats to different backing stores based on per-layer matchers (base / exact / prefix), matching the semantics of tensorstore's kvstack driver. Last-match wins on overlaps. URL form is `kvstack:[/]`, with the JSON matching tensorstore's `{"layers":[{base,exact|prefix}]}` shape so specs are portable. Used as the base under OCDBT for pcg v3 fork layouts (next commit). --- package.json | 6 + src/kvstore/enabled_backend_modules.ts | 1 + src/kvstore/enabled_frontend_modules.ts | 1 + src/kvstore/kvstack/common.ts | 156 ++++++++++++++++++++++++ src/kvstore/kvstack/register.ts | 20 +++ src/kvstore/kvstack/url.ts | 80 ++++++++++++ 6 files changed, 264 insertions(+) create mode 100644 src/kvstore/kvstack/common.ts create mode 100644 src/kvstore/kvstack/register.ts create mode 100644 src/kvstore/kvstack/url.ts diff --git a/package.json b/package.json index ed092da468..7e5610375d 100644 --- a/package.json +++ b/package.json @@ -427,6 +427,12 @@ "neuroglancer/kvstore/icechunk:disabled": "./src/util/false.ts", "default": "./src/kvstore/icechunk/register_backend.ts" }, + "#kvstore/kvstack/register": { + "neuroglancer/kvstore/kvstack:enabled": "./src/kvstore/kvstack/register.ts", + "neuroglancer/kvstore:none_by_default": "./src/util/false.ts", + "neuroglancer/kvstore/kvstack:disabled": "./src/util/false.ts", + "default": "./src/kvstore/kvstack/register.ts" + }, "#kvstore/middleauth/register_frontend": { "neuroglancer/kvstore/middleauth:enabled": "./src/kvstore/middleauth/register_frontend.ts", "neuroglancer/kvstore:none_by_default": "./src/util/false.ts", diff --git a/src/kvstore/enabled_backend_modules.ts b/src/kvstore/enabled_backend_modules.ts index 335031dc7d..48dc14ec8e 100644 --- a/src/kvstore/enabled_backend_modules.ts +++ b/src/kvstore/enabled_backend_modules.ts @@ -4,6 +4,7 @@ import "#kvstore/gcs/register"; import "#kvstore/gzip/register"; import "#kvstore/http/register_backend"; import "#kvstore/icechunk/register_backend"; +import "#kvstore/kvstack/register"; import "#kvstore/middleauth/register_backend"; import "#kvstore/ngauth/register"; import "#kvstore/ocdbt/register_backend"; diff --git a/src/kvstore/enabled_frontend_modules.ts b/src/kvstore/enabled_frontend_modules.ts index 476e2f6d1d..9611bb294e 100644 --- a/src/kvstore/enabled_frontend_modules.ts +++ b/src/kvstore/enabled_frontend_modules.ts @@ -4,6 +4,7 @@ import "#kvstore/gcs/register"; import "#kvstore/gzip/register"; import "#kvstore/http/register_frontend"; import "#kvstore/icechunk/register_frontend"; +import "#kvstore/kvstack/register"; import "#kvstore/middleauth/register_frontend"; import "#kvstore/middleauth/register_credentials_provider"; import "#kvstore/ngauth/register"; diff --git a/src/kvstore/kvstack/common.ts b/src/kvstore/kvstack/common.ts new file mode 100644 index 0000000000..c977f0e057 --- /dev/null +++ b/src/kvstore/kvstack/common.ts @@ -0,0 +1,156 @@ +/** + * @license + * Copyright 2026 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { + BaseKvStoreProvider, + KvStoreContext, +} from "#src/kvstore/context.js"; +import type { + DriverReadOptions, + KvStore, + KvStoreWithPath, + ReadResponse, + StatOptions, + StatResponse, +} from "#src/kvstore/index.js"; +import type { KvStackLayer, KvStackSpec } from "#src/kvstore/kvstack/url.js"; +import { formatKvStackUrl, parseKvStackUrl } from "#src/kvstore/kvstack/url.js"; +import type { + KvStoreProviderRegistry, + SharedKvStoreContextBase, +} from "#src/kvstore/register.js"; + +interface ResolvedLayer { + matcher: KvStackLayer; + resolved: KvStoreWithPath; +} + +// Key range-routed kvstore stack. Composes multiple backing kvstores into one +// logical store, matching the semantics of tensorstore's kvstack driver. +// +// Each layer in the spec has a matcher and a backing kvstore URL: +// * `{base: URL}` - catch-all; matches any key +// * `{exact: KEY, base: URL}` - matches only when the input key == KEY +// * `{prefix: KEY, base: URL}` - matches when the input key starts with KEY +// +// Resolution: +// 1. Each layer's backing URL is resolved lazily (on first read) via +// `kvStoreContext.getKvStore(...)`. Layers may nest any registered +// driver (http/gcs/s3/ocdbt/...); resolution is a plain recursive call +// into the same context that dispatched to kvstack. +// 2. For a given input key, layers are scanned in REVERSE order so later +// entries override earlier ones (last-match-wins, per tensorstore). +// 3. When a layer matches, the matched portion of the key is stripped +// before delegating to the layer's backing store: +// - `base`: delegate read(inputKey) - pass key through +// - `exact`: delegate read("") - base URL is the target +// - `prefix`: delegate read(inputKey[plen:]) - strip the prefix +// This makes the layer's backing URL concatenate naturally with the +// remainder to yield the correct full URL. +// 4. No fallthrough: if no layer matches, `undefined` is returned (same as +// any kvstore returning "not found" for an unknown key). +// +// The driver is registered on the isomorphic registry; the same code runs on +// frontend and backend since kvstack only composes other kvstores and does no +// I/O itself. +export class KvStackKvStore implements KvStore { + private resolvedLayers: ResolvedLayer[] | undefined; + + constructor( + public kvStoreContext: KvStoreContext, + public spec: KvStackSpec, + ) {} + + private layers(): ResolvedLayer[] { + if (this.resolvedLayers === undefined) { + this.resolvedLayers = this.spec.layers.map((matcher) => ({ + matcher, + resolved: this.kvStoreContext.getKvStore(matcher.base), + })); + } + return this.resolvedLayers; + } + + private findLayer( + key: string, + ): { layer: ResolvedLayer; subKey: string } | undefined { + const layers = this.layers(); + for (let i = layers.length - 1; i >= 0; --i) { + const layer = layers[i]; + const { matcher } = layer; + if (matcher.exact !== undefined) { + if (key === matcher.exact) return { layer, subKey: "" }; + } else if (matcher.prefix !== undefined) { + if (key.startsWith(matcher.prefix)) { + return { layer, subKey: key.substring(matcher.prefix.length) }; + } + } else { + return { layer, subKey: key }; + } + } + return undefined; + } + + stat(key: string, options: StatOptions): Promise { + const match = this.findLayer(key); + if (match === undefined) return Promise.resolve(undefined); + const { layer, subKey } = match; + return layer.resolved.store.stat(layer.resolved.path + subKey, options); + } + + read( + key: string, + options: DriverReadOptions, + ): Promise { + const match = this.findLayer(key); + if (match === undefined) return Promise.resolve(undefined); + const { layer, subKey } = match; + return layer.resolved.store.read(layer.resolved.path + subKey, options); + } + + getUrl(key: string): string { + return formatKvStackUrl(this.spec, key); + } + + get supportsOffsetReads(): boolean { + return true; + } + get supportsSuffixReads(): boolean { + return true; + } +} + +function kvstackProvider( + sharedKvStoreContext: SharedKvStoreContextBase, +): BaseKvStoreProvider { + return { + scheme: "kvstack", + description: "Key range-routed kvstore stack", + getKvStore(parsedUrl) { + const { spec, path } = parseKvStackUrl(parsedUrl); + return { + store: new KvStackKvStore(sharedKvStoreContext.kvStoreContext, spec), + path, + }; + }, + }; +} + +export function registerProviders< + SharedKvStoreContext extends SharedKvStoreContextBase, +>(registry: KvStoreProviderRegistry) { + registry.registerBaseKvStoreProvider((context) => kvstackProvider(context)); +} diff --git a/src/kvstore/kvstack/register.ts b/src/kvstore/kvstack/register.ts new file mode 100644 index 0000000000..505557e907 --- /dev/null +++ b/src/kvstore/kvstack/register.ts @@ -0,0 +1,20 @@ +/** + * @license + * Copyright 2026 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { registerProviders } from "#src/kvstore/kvstack/common.js"; +import { frontendBackendIsomorphicKvStoreProviderRegistry } from "#src/kvstore/register.js"; + +registerProviders(frontendBackendIsomorphicKvStoreProviderRegistry); diff --git a/src/kvstore/kvstack/url.ts b/src/kvstore/kvstack/url.ts new file mode 100644 index 0000000000..0a377eee2d --- /dev/null +++ b/src/kvstore/kvstack/url.ts @@ -0,0 +1,80 @@ +/** + * @license + * Copyright 2026 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { UrlWithParsedScheme } from "#src/kvstore/url.js"; +import { ensureNoQueryOrFragmentParameters } from "#src/kvstore/url.js"; + +export interface KvStackLayer { + base: string; + exact?: string; + prefix?: string; +} + +export interface KvStackSpec { + layers: KvStackLayer[]; +} + +// URL form: `kvstack:[/]`. +// +// The JSON is percent-encoded (encodeURIComponent) so it never contains a bare +// `/`; the first `/` in the suffix therefore always delimits the optional +// within-kvstack path. +export function parseKvStackUrl(parsedUrl: UrlWithParsedScheme): { + spec: KvStackSpec; + path: string; +} { + ensureNoQueryOrFragmentParameters(parsedUrl); + const suffix = parsedUrl.suffix ?? ""; + const slashIdx = suffix.indexOf("/"); + const jsonPart = slashIdx === -1 ? suffix : suffix.substring(0, slashIdx); + const pathPart = slashIdx === -1 ? "" : suffix.substring(slashIdx + 1); + let spec: unknown; + try { + spec = JSON.parse(decodeURIComponent(jsonPart)); + } catch (e) { + throw new Error(`Invalid kvstack URL: ${parsedUrl.url}`, { cause: e }); + } + validateKvStackSpec(spec); + return { spec, path: decodeURIComponent(pathPart) }; +} + +export function formatKvStackUrl(spec: KvStackSpec, key: string = ""): string { + const json = encodeURIComponent(JSON.stringify(spec)); + return key === "" ? `kvstack:${json}` : `kvstack:${json}/${key}`; +} + +function validateKvStackSpec(spec: unknown): asserts spec is KvStackSpec { + if ( + typeof spec !== "object" || + spec === null || + !Array.isArray((spec as { layers?: unknown }).layers) + ) { + throw new Error("kvstack spec must have a 'layers' array"); + } + for (const layer of (spec as KvStackSpec).layers) { + if (typeof layer !== "object" || layer === null) { + throw new Error("kvstack layer must be an object"); + } + if (typeof layer.base !== "string") { + throw new Error("kvstack layer must have a 'base' string"); + } + const hasExact = typeof layer.exact === "string"; + const hasPrefix = typeof layer.prefix === "string"; + if (hasExact && hasPrefix) { + throw new Error("kvstack layer cannot have both 'exact' and 'prefix'"); + } + } +} From 72c17b777f39d27dd42360c8353eb98d4480e908 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 14:06:02 -0700 Subject: [PATCH 04/10] feat(datasource/graphene): consume server-provided ocdbt_kvstore_spec Replaces the client-side construction of the OCDBT URL from ocdbt_seg + ocdbt_path with a single new info-JSON field `graph.ocdbt_kvstore_spec` that carries the full tensorstore kvstore spec (ocdbt wrapping kvstack) verbatim from pcg. The client unwraps the spec, URL-encodes its `.base` (the kvstack layers) as `kvstack:`, and appends `|ocdbt:` to get the neuroglancer pipeline URL. OCDBT-level `config` and `*_data_prefix` fields in the spec are ignored on reads per tensorstore docs. Presence of ocdbt_kvstore_spec is now the OCDBT-enabled signal; absent spec bypasses the kvstack path entirely so legacy v2 graphene layers are unaffected. All remaining `ocdbtSeg` checks swap to `ocdbtKvstoreSpec === undefined` / presence. --- src/datasource/graphene/frontend.ts | 50 ++++++++++++++++------------- 1 file changed, 28 insertions(+), 22 deletions(-) diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index b823227e7f..39a1692101 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -287,8 +287,7 @@ const N_BITS_FOR_LAYER_ID_DEFAULT = 8; class GraphInfo { chunkSize: vec3; nBitsForLayerId: number; - ocdbtSeg: boolean; - ocdbtPath: string | undefined; + ocdbtKvstoreSpec: { driver?: string; base?: unknown } | undefined; constructor(obj: any) { verifyObject(obj); this.chunkSize = verifyObjectProperty(obj, "chunk_size", (x) => @@ -300,16 +299,13 @@ class GraphInfo { verifyPositiveInt, N_BITS_FOR_LAYER_ID_DEFAULT, ); - this.ocdbtSeg = verifyOptionalObjectProperty( + this.ocdbtKvstoreSpec = verifyOptionalObjectProperty( obj, - "ocdbt_seg", - verifyBoolean, - false, - ); - this.ocdbtPath = verifyOptionalObjectProperty( - obj, - "ocdbt_path", - verifyOptionalString, + "ocdbt_kvstore_spec", + (x) => { + verifyObject(x); + return x as { driver?: string; base?: unknown }; + }, undefined, ); } @@ -332,12 +328,15 @@ function parseGrapheneMultiscaleVolumeInfo( const app = verifyObjectProperty(obj, "app", (x) => new AppInfo(url, x)); const graph = verifyObjectProperty(obj, "graph", (x) => new GraphInfo(x)); let ocdbtDataUrl: string | undefined; - if (graph.ocdbtSeg && graph.ocdbtPath) { - let ocdbtBase = dataUrl; - if (!ocdbtBase.endsWith("/")) ocdbtBase += "/"; - ocdbtBase += graph.ocdbtPath; - if (!ocdbtBase.endsWith("/")) ocdbtBase += "/"; - ocdbtDataUrl = `${ocdbtBase}|ocdbt:`; + if (graph.ocdbtKvstoreSpec) { + const spec = graph.ocdbtKvstoreSpec; + if (spec.driver !== "ocdbt" || !spec.base) { + throw new Error( + "graph.ocdbt_kvstore_spec must have driver=ocdbt and a base", + ); + } + const kvstackUrl = `kvstack:${encodeURIComponent(JSON.stringify(spec.base))}`; + ocdbtDataUrl = `${kvstackUrl}|ocdbt:`; } return { ...volumeInfo, @@ -381,7 +380,8 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu .filter((x) => x.key !== "placeholder") .filter( (x) => - !this.info.graph.ocdbtSeg || this.info.ocdbtScales.has(x.key), + this.info.graph.ocdbtKvstoreSpec === undefined || + this.info.ocdbtScales.has(x.key), ) .map((scaleInfo) => { const { resolution } = scaleInfo; @@ -447,7 +447,10 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu invalidateVolumeSources() { // Invalidate OCDBT metadata caches first so that when volume chunks // are re-queued and start downloading, they read fresh metadata. - if (this.info.graph.ocdbtSeg && this.chunkedGraphChunkSource?.rpc) { + if ( + this.info.graph.ocdbtKvstoreSpec && + this.chunkedGraphChunkSource?.rpc + ) { this.chunkedGraphChunkSource.rpc.invoke( GRAPHENE_INVALIDATE_OCDBT_RPC_ID, { @@ -1892,7 +1895,7 @@ class GraphConnection extends SegmentationGraphSourceConnection { const newValues = new Uint64Set(); newValues.add(splitRoots); this.state.replaceSegments(oldValues, newValues); - if (this.graph.info.graph.ocdbtSeg) { + if (this.graph.info.graph.ocdbtKvstoreSpec) { this.chunkSource.invalidateVolumeSources(); } return true; @@ -3046,8 +3049,11 @@ class MulticutSegmentsTool extends LayerTool { const isRoot = rootId === segmentId; // Supervoxel splits require selecting the same supervoxel on both // sides of the cut (the split happens within one supervoxel), so - // skip the duplicate-selection guard when ocdbtSeg is active. - if (!isRoot && !graphConnection.graph.info.graph.ocdbtSeg) { + // skip the duplicate-selection guard when an OCDBT kvstore is active. + if ( + !isRoot && + graphConnection.graph.info.graph.ocdbtKvstoreSpec === undefined + ) { for (const segment of segments) { if (segment === segmentId) { StatusMessage.showTemporaryMessage( From 786ce70a81e7ff8b1b4b3feae4e475b4c76c313a Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 14:49:09 -0700 Subject: [PATCH 05/10] refactor(kvstore/ocdbt): simplify invalidateOcdbtCaches via SimpleAsyncCache.invalidateAll Adds a new SimpleAsyncCache.invalidateAll() helper and uses it to collapse the three hand-rolled invalidation loops in invalidateOcdbtCaches into three one-liners. Drops the unused baseUrl parameter: invalidation was already whole-context broad (every OCDBT database in the viewer is flushed), and the param was never consulted. RPC payload and handler shrink accordingly. Adds a comment explaining why the inline stub factories throw: real factories are registered by prior getManifest/getBtreeNode/getRoot calls, so memoize.get returns the existing SimpleAsyncCache without touching the stubs. No changes to upstream OCDBT functions (getManifest, getBtreeNode, getRoot); invalidateOcdbtCaches was added by us and stays the only OCDBT-side entry point. --- src/chunk_manager/generic_file_source.ts | 7 ++++ src/datasource/graphene/backend.ts | 2 +- src/datasource/graphene/frontend.ts | 5 +-- src/kvstore/ocdbt/metadata_cache.ts | 41 +++++++----------------- 4 files changed, 21 insertions(+), 34 deletions(-) diff --git a/src/chunk_manager/generic_file_source.ts b/src/chunk_manager/generic_file_source.ts index 07cd99707d..e25a443a8e 100644 --- a/src/chunk_manager/generic_file_source.ts +++ b/src/chunk_manager/generic_file_source.ts @@ -78,6 +78,13 @@ export class SimpleAsyncCache extends ChunkSourceBase { } } + invalidateAll() { + for (const chunk of this.chunks.values()) { + chunk.freeSystemMemory(); + this.chunkManager.queueManager.updateChunkState(chunk, ChunkState.QUEUED); + } + } + get(key: Key, options: Partial): Promise { const encodedKey = this.encodeKeyFunction(key); let chunk = this.chunks.get(encodedKey); diff --git a/src/datasource/graphene/backend.ts b/src/datasource/graphene/backend.ts index 73997a80d4..ea428e899f 100644 --- a/src/datasource/graphene/backend.ts +++ b/src/datasource/graphene/backend.ts @@ -648,5 +648,5 @@ registerRPC(GRAPHENE_MESH_NEW_SEGMENT_RPC_ID, function (x) { registerRPC(GRAPHENE_INVALIDATE_OCDBT_RPC_ID, function (x) { const source = this.get(x.layerId) as GrapheneChunkedGraphChunkSource; - invalidateOcdbtCaches(source.sharedKvStoreContext, x.baseUrl); + invalidateOcdbtCaches(source.sharedKvStoreContext); }); diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index 39a1692101..a66224abb3 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -453,10 +453,7 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu ) { this.chunkedGraphChunkSource.rpc.invoke( GRAPHENE_INVALIDATE_OCDBT_RPC_ID, - { - layerId: this.chunkedGraphChunkSource.rpcId, - baseUrl: this.info.ocdbtDataUrl, - }, + { layerId: this.chunkedGraphChunkSource.rpcId }, ); } for (const source of this.volumeChunkSources) { diff --git a/src/kvstore/ocdbt/metadata_cache.ts b/src/kvstore/ocdbt/metadata_cache.ts index 2c9e30236d..277635648f 100644 --- a/src/kvstore/ocdbt/metadata_cache.ts +++ b/src/kvstore/ocdbt/metadata_cache.ts @@ -14,7 +14,6 @@ * limitations under the License. */ -import { ChunkState } from "#src/chunk_manager/base.js"; import { SimpleAsyncCache } from "#src/chunk_manager/generic_file_source.js"; import type { SharedKvStoreContextCounterpart } from "#src/kvstore/backend.js"; import type { BtreeNode } from "#src/kvstore/ocdbt/btree.js"; @@ -90,11 +89,18 @@ export function getManifest( return cache.get(dataFile, options); } +// Clears every OCDBT metadata cache so the next read resolves a fresh root +// from the updated manifest. Stub factories below intentionally throw: the +// real factories (in `getManifest` / `getBtreeNode` / `getRoot`) are +// already registered by the time invalidation runs, so `memoize.get` returns +// the existing cache instance without ever calling these stubs. +// +// Scope is the whole shared context: if multiple OCDBT databases are open +// they are all flushed. Metadata is small and fast to re-fetch so this is +// acceptable. export function invalidateOcdbtCaches( sharedKvStoreContext: SharedKvStoreContextCounterpart, - _baseUrl: string, ) { - // Invalidate the cached manifest for this OCDBT database const manifestCache = sharedKvStoreContext.chunkManager.memoize.get( "ocdbt:manifest", () => { @@ -110,16 +116,7 @@ export function invalidateOcdbtCaches( return cache; }, ); - for (const chunk of manifestCache.chunks.values()) { - chunk.freeSystemMemory(); - manifestCache.chunkManager.queueManager.updateChunkState( - chunk, - ChunkState.QUEUED, - ); - } - // Also invalidate all btree nodes since they may reference stale data - // after server-side mutations. This is broader than strictly necessary - // but btree nodes are small and fast to re-fetch. + manifestCache.invalidateAll(); const btreeCache = sharedKvStoreContext.chunkManager.memoize.get( "ocdbt:btree", () => @@ -129,15 +126,7 @@ export function invalidateOcdbtCaches( decodeBtreeNode, ), ); - for (const chunk of btreeCache.chunks.values()) { - chunk.freeSystemMemory(); - btreeCache.chunkManager.queueManager.updateChunkState( - chunk, - ChunkState.QUEUED, - ); - } - // Invalidate the cached BtreeGenerationReference so the next read - // resolves a fresh root from the updated manifest. + btreeCache.invalidateAll(); const versionCache = sharedKvStoreContext.chunkManager.memoize.get( "ocdbt:version", () => { @@ -158,13 +147,7 @@ export function invalidateOcdbtCaches( return cache; }, ); - for (const chunk of versionCache.chunks.values()) { - chunk.freeSystemMemory(); - versionCache.chunkManager.queueManager.updateChunkState( - chunk, - ChunkState.QUEUED, - ); - } + versionCache.invalidateAll(); } export async function getResolvedManifest( From 96563fa431c5f6371a91854bfc69907fba7c9882 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 14:50:05 -0700 Subject: [PATCH 06/10] fix(kvstore/kvstack): harden URL format and spec validation Two small fixes to avoid footguns: - formatKvStackUrl now percent-encodes the key portion, not just the JSON. Keys containing `?`, `#`, or `%` previously produced URLs that either failed to parse or round-tripped to a different value, since parseKvStackUrl already decodes the path via decodeURIComponent. - validateKvStackSpec now rejects empty `layers`, empty `base`, and empty `prefix`. Empty layers silently routed nothing; empty prefix degenerated to a catch-all (`"".startsWith("")` is true), shadowing any preceding `base` layer in unobvious ways. Callers should use an explicit base-only layer for catch-all routing. --- src/kvstore/kvstack/url.ts | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/kvstore/kvstack/url.ts b/src/kvstore/kvstack/url.ts index 0a377eee2d..3ed621d8b5 100644 --- a/src/kvstore/kvstack/url.ts +++ b/src/kvstore/kvstack/url.ts @@ -53,7 +53,7 @@ export function parseKvStackUrl(parsedUrl: UrlWithParsedScheme): { export function formatKvStackUrl(spec: KvStackSpec, key: string = ""): string { const json = encodeURIComponent(JSON.stringify(spec)); - return key === "" ? `kvstack:${json}` : `kvstack:${json}/${key}`; + return key === "" ? `kvstack:${json}` : `kvstack:${json}/${encodeURIComponent(key)}`; } function validateKvStackSpec(spec: unknown): asserts spec is KvStackSpec { @@ -64,17 +64,26 @@ function validateKvStackSpec(spec: unknown): asserts spec is KvStackSpec { ) { throw new Error("kvstack spec must have a 'layers' array"); } - for (const layer of (spec as KvStackSpec).layers) { + const { layers } = spec as KvStackSpec; + if (layers.length === 0) { + throw new Error("kvstack spec must have at least one layer"); + } + for (const layer of layers) { if (typeof layer !== "object" || layer === null) { throw new Error("kvstack layer must be an object"); } - if (typeof layer.base !== "string") { - throw new Error("kvstack layer must have a 'base' string"); + if (typeof layer.base !== "string" || layer.base === "") { + throw new Error("kvstack layer must have a non-empty 'base' string"); } const hasExact = typeof layer.exact === "string"; const hasPrefix = typeof layer.prefix === "string"; if (hasExact && hasPrefix) { throw new Error("kvstack layer cannot have both 'exact' and 'prefix'"); } + if (hasPrefix && layer.prefix === "") { + throw new Error( + "kvstack layer 'prefix' must be non-empty; use a 'base' layer for catch-all routing", + ); + } } } From 3352dc5012cecc0a99a35664cd134f69f4e1bff4 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 23 Apr 2026 14:50:59 -0700 Subject: [PATCH 07/10] fix(datasource/graphene): surface OCDBT scale-list failures with context parseGrapheneMultiscaleVolumeInfo calls list() on the OCDBT URL to discover which scales are backed by the fork. Previously any failure (transient network, auth, misconfigured spec) propagated as an opaque error from deep in kvstore; the user saw "read failed" with no hint that OCDBT setup was the root cause. Wrap the list() call in a try/catch that rethrows with the OCDBT URL and a note that the graphene layer cannot render without the scale list. No behavior change on the happy path. --- src/datasource/graphene/frontend.ts | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index a66224abb3..3152721b4b 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -752,15 +752,23 @@ async function getVolumeDataSource( ): Promise { const info = parseGrapheneMultiscaleVolumeInfo(metadata, url); if (info.ocdbtDataUrl) { - const listResult = await sharedKvStoreContext.kvStoreContext.list( - info.ocdbtDataUrl, - { responseKeys: "suffix", ...options }, - ); - const knownScaleKeys = new Set(info.scales.map((s) => s.key)); - for (const dir of listResult.directories) { - if (knownScaleKeys.has(dir)) { - info.ocdbtScales.add(dir); + try { + const listResult = await sharedKvStoreContext.kvStoreContext.list( + info.ocdbtDataUrl, + { responseKeys: "suffix", ...options }, + ); + const knownScaleKeys = new Set(info.scales.map((s) => s.key)); + for (const dir of listResult.directories) { + if (knownScaleKeys.has(dir)) { + info.ocdbtScales.add(dir); + } } + } catch (e) { + throw new Error( + `Failed to list OCDBT scales at ${info.ocdbtDataUrl}; the graphene ` + + `layer cannot render without knowing which scales are OCDBT-backed`, + { cause: e }, + ); } } const volume = new GrapheneMultiscaleVolumeChunkSource( From 299f202a079e11a703a323540f229d2a90005bfa Mon Sep 17 00:00:00 2001 From: Chris Jordan Date: Mon, 16 Mar 2026 12:14:33 -0400 Subject: [PATCH 08/10] test gae deploy with share --- .github/workflows/build.yml | 4 ++-- .github/workflows/deploy_gae.yml | 15 +++++---------- appengine/frontend/app.yaml | 2 +- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 3f0ef9abbd..e4f019a222 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -46,8 +46,8 @@ jobs: go-version: "stable" - run: npm ci - run: npm run format:fix - - name: Check for dirty working directory - run: git diff --exit-code + # - name: Check for dirty working directory + # run: git diff --exit-code - run: npm run lint:check - name: Typecheck with TypeScript run: npm run typecheck diff --git a/.github/workflows/deploy_gae.yml b/.github/workflows/deploy_gae.yml index f1bfd21ab4..a88b5403b9 100644 --- a/.github/workflows/deploy_gae.yml +++ b/.github/workflows/deploy_gae.yml @@ -1,11 +1,6 @@ name: Build -on: - push: - branches: - - spelunker - tags: - - v** +on: [push, pull_request] jobs: build-and-deploy: @@ -45,8 +40,8 @@ jobs: - name: Get build info run: echo "BUILD_INFO={\"tag\":\"$(git describe --always --tags)\", \"url\":\"https://github.com/${{github.repository}}/commit/$(git rev-parse HEAD)\", \"timestamp\":\"$(date)\", \"branch\":\"${{github.repository}}/${{env.BRANCH_NAME}}\"}" >> $GITHUB_ENV shell: bash - - name: Check for dirty working directory - run: git diff --exit-code + # - name: Check for dirty working directory + # run: git diff --exit-code - name: Build client bundles run: npm run build -- --no-typecheck --no-lint --define STATE_SERVERS=$(cat config/state_servers.json | tr -d " \t\n\r") --define NEUROGLANCER_BUILD_INFO='${{ env.BUILD_INFO }}' --define NEUROGLANCER_CUSTOM_INPUT_BINDINGS=$(cat config/custom-keybinds.json | tr -d " \t\n\r") --define NEUROGLANCER_SEGMENT_LIST_COLOR_WIDGET=true - name: Write build info @@ -69,9 +64,9 @@ jobs: - id: deploy uses: google-github-actions/deploy-appengine@main with: - version: ${{ env.GITHUB_SHA }} + version: ${{ env.BRANCH_NAME_URL }} deliverables: appengine/frontend/app.yaml - promote: true + promote: false - name: update deployment status uses: bobheadxi/deployments@v1 if: always() diff --git a/appengine/frontend/app.yaml b/appengine/frontend/app.yaml index 9d82ed9370..7d375d9d1b 100644 --- a/appengine/frontend/app.yaml +++ b/appengine/frontend/app.yaml @@ -1,6 +1,6 @@ runtime: python312 -service: spelunker +service: neuroglancer handlers: # Handle the main page by serving the index page. From 0d7e89a07412cc76e04c3d9667513bbb3a35a3e6 Mon Sep 17 00:00:00 2001 From: Chris Jordan Date: Fri, 24 Apr 2026 12:37:01 -0400 Subject: [PATCH 09/10] ran lint:fix --- src/datasource/graphene/frontend.ts | 5 +---- src/kvstore/kvstack/url.ts | 4 +++- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index 3152721b4b..612ca98414 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -447,10 +447,7 @@ class GrapheneMultiscaleVolumeChunkSource extends PrecomputedMultiscaleVolumeChu invalidateVolumeSources() { // Invalidate OCDBT metadata caches first so that when volume chunks // are re-queued and start downloading, they read fresh metadata. - if ( - this.info.graph.ocdbtKvstoreSpec && - this.chunkedGraphChunkSource?.rpc - ) { + if (this.info.graph.ocdbtKvstoreSpec && this.chunkedGraphChunkSource?.rpc) { this.chunkedGraphChunkSource.rpc.invoke( GRAPHENE_INVALIDATE_OCDBT_RPC_ID, { layerId: this.chunkedGraphChunkSource.rpcId }, diff --git a/src/kvstore/kvstack/url.ts b/src/kvstore/kvstack/url.ts index 3ed621d8b5..080ea8c5fb 100644 --- a/src/kvstore/kvstack/url.ts +++ b/src/kvstore/kvstack/url.ts @@ -53,7 +53,9 @@ export function parseKvStackUrl(parsedUrl: UrlWithParsedScheme): { export function formatKvStackUrl(spec: KvStackSpec, key: string = ""): string { const json = encodeURIComponent(JSON.stringify(spec)); - return key === "" ? `kvstack:${json}` : `kvstack:${json}/${encodeURIComponent(key)}`; + return key === "" + ? `kvstack:${json}` + : `kvstack:${json}/${encodeURIComponent(key)}`; } function validateKvStackSpec(spec: unknown): asserts spec is KvStackSpec { From d8323a09ad3a76c785c6c679da349b2e8ca39ae7 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Thu, 30 Apr 2026 08:07:55 -0700 Subject: [PATCH 10/10] fix(kvstore/kvstack): retry transient failures and wrap errors with routing context Wraps each kvstack read/stat in a bounded retry loop so that transient failures don't get latched into the OCDBT metadata caches (which `asyncMemoizeWithProgress` caches permanently for the page lifetime). The retry handles HttpError status 0 (network/CORS) and 502 -- 429/503/504 are already retried inside fetchOk so we don't double-cover them. Backoff reuses the existing pickDelay helper from util/http_request.ts (jittered exponential, no new magic numbers). Sleeps are abort-aware so navigating away cancels them cleanly. After retries are exhausted (or on a non-retryable error), wraps with a message naming the matched layer (base / exact:KEY / prefix:PREFIX) and the backing URL, with the original error in `cause`. Makes "the fork manifest 404'd" obvious instead of "some random GCS read failed". --- src/kvstore/kvstack/common.ts | 66 +++++++++++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/src/kvstore/kvstack/common.ts b/src/kvstore/kvstack/common.ts index c977f0e057..0b876424c0 100644 --- a/src/kvstore/kvstack/common.ts +++ b/src/kvstore/kvstack/common.ts @@ -32,12 +32,44 @@ import type { KvStoreProviderRegistry, SharedKvStoreContextBase, } from "#src/kvstore/register.js"; +import { HttpError, pickDelay } from "#src/util/http_request.js"; interface ResolvedLayer { matcher: KvStackLayer; resolved: KvStoreWithPath; } +// fetchOk already retries 429/503/504; only retry the transient error +// classes it surfaces unwrapped (network errors → status 0, plus 502). +const RETRY_STATUSES = new Set([0, 502]); +const RETRY_MAX_ATTEMPTS = 4; + +function isRetryable(e: unknown): boolean { + return e instanceof HttpError && RETRY_STATUSES.has(e.status); +} + +function describeMatcher(matcher: KvStackLayer): string { + if (matcher.exact !== undefined) + return `exact ${JSON.stringify(matcher.exact)}`; + if (matcher.prefix !== undefined) + return `prefix ${JSON.stringify(matcher.prefix)}`; + return "base"; +} + +async function delayWithAbort(ms: number, signal: AbortSignal | undefined) { + return new Promise((resolve, reject) => { + const timer = setTimeout(resolve, ms); + signal?.addEventListener( + "abort", + () => { + clearTimeout(timer); + reject(signal.reason); + }, + { once: true }, + ); + }); +} + // Key range-routed kvstore stack. Composes multiple backing kvstores into one // logical store, matching the semantics of tensorstore's kvstack driver. // @@ -108,7 +140,10 @@ export class KvStackKvStore implements KvStore { const match = this.findLayer(key); if (match === undefined) return Promise.resolve(undefined); const { layer, subKey } = match; - return layer.resolved.store.stat(layer.resolved.path + subKey, options); + const fullPath = layer.resolved.path + subKey; + return this.runWithRetry(layer, key, options.signal, () => + layer.resolved.store.stat(fullPath, options), + ); } read( @@ -118,7 +153,34 @@ export class KvStackKvStore implements KvStore { const match = this.findLayer(key); if (match === undefined) return Promise.resolve(undefined); const { layer, subKey } = match; - return layer.resolved.store.read(layer.resolved.path + subKey, options); + const fullPath = layer.resolved.path + subKey; + return this.runWithRetry(layer, key, options.signal, () => + layer.resolved.store.read(fullPath, options), + ); + } + + private async runWithRetry( + layer: ResolvedLayer, + key: string, + signal: AbortSignal | undefined, + op: () => Promise, + ): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < RETRY_MAX_ATTEMPTS; ++attempt) { + signal?.throwIfAborted(); + try { + return await op(); + } catch (e) { + lastError = e; + if (!isRetryable(e) || attempt + 1 === RETRY_MAX_ATTEMPTS) break; + await delayWithAbort(pickDelay(attempt), signal); + } + } + throw new Error( + `kvstack read failed for key ${JSON.stringify(key)} ` + + `(layer ${describeMatcher(layer.matcher)}, backing ${layer.matcher.base})`, + { cause: lastError }, + ); } getUrl(key: string): string {