fix: Bring back refresh request after reroot

src/datasource/catmaid/api.ts

@@ -1183,30 +1183,6 @@ function isCatmaidServerVersionSupported(version: string | undefined) {
   );
 }
 
-function makeCatmaidNodeRevisionUpdates(
-  nodes: readonly CatmaidEditParentContext[] | undefined,
-  revisionToken: string,
-): readonly CatmaidSkeletonNodeSourceStateUpdate[] {
-  if (nodes === undefined) {
-    return [];
-  }
-  const sourceState = { revisionToken };
-  const seen = new Set<number>();
-  const revisionUpdates: CatmaidSkeletonNodeSourceStateUpdate[] = [];
-  for (const node of nodes) {
-    const nodeId = Number(node.nodeId);
-    if (!Number.isFinite(nodeId)) continue;
-    const normalizedNodeId = Math.round(nodeId);
-    if (seen.has(normalizedNodeId)) continue;
-    seen.add(normalizedNodeId);
-    revisionUpdates.push({
-      nodeId: normalizedNodeId,
-      sourceState,
-    });
-  }
-  return revisionUpdates;
-}
-
 function fetchWithCatmaidCredentials(
   credentialsProvider: CredentialsProvider<CatmaidToken>,
   input: string,
@@ -1691,18 +1667,12 @@ export class CatmaidClient implements CatmaidSpatialSkeletonEditApi {
       method: "POST",
       body,
     });
-    const revisionToken = normalizeCatmaidRevisionToken(response?.edition_time);
-    if (revisionToken === undefined) {
+    if (Number(response?.newroot) !== nodeId) {
       throw new Error(
-        "CATMAID skeleton/reroot did not return the new root edition_time.",
+        "CATMAID skeleton/reroot did not return the requested new root.",
       );
     }
-    return {
-      nodeSourceStateUpdates: makeCatmaidNodeRevisionUpdates(
-        editContext?.nodes,
-        revisionToken,
-      ),
-    };
+    return {};
   }
 
   async deleteNode(

src/datasource/catmaid/spatial_skeleton_commands.ts

@@ -553,6 +553,7 @@ interface ResolvedSpatialSkeletonEditNodeContext {
 }
 
 type CatmaidSkeletonRootNodeSource = Pick<CatmaidClient, "getSkeletonRootNode">;
+type CatmaidSkeletonSourceStateRefresh = Pick<CatmaidClient, "getSkeleton">;
 
 function collectUniqueNodePositions(
   ...nodeSets: readonly (readonly (
@@ -639,6 +640,42 @@ function getCatmaidSkeletonRootNodeSource(
     : undefined;
 }
 
+async function getFreshRerootSourceStateUpdates(
+  skeletonSource: CatmaidSkeletonSourceStateRefresh,
+  segmentId: number,
+  nodeIds: readonly number[],
+): Promise<readonly CatmaidSpatialSkeletonNodeSourceStateUpdate[]> {
+  const refreshedNodes = await skeletonSource.getSkeleton(segmentId);
+  const refreshedNodeById = new Map(
+    refreshedNodes.map((node) => [node.nodeId, node]),
+  );
+  const seen = new Set<number>();
+  const updates: CatmaidSpatialSkeletonNodeSourceStateUpdate[] = [];
+  for (const nodeId of nodeIds) {
+    if (seen.has(nodeId)) continue;
+    seen.add(nodeId);
+    const sourceState = refreshedNodeById.get(nodeId)?.sourceState;
+    if (sourceState === undefined) {
+      throw new Error(
+        `CATMAID reroot refresh did not return revision state for node ${nodeId}.`,
+      );
+    }
+    updates.push({ nodeId, sourceState });
+  }
+  return updates;
+}
+
+class CatmaidRerootSourceStateRefreshError extends Error {
+  constructor(readonly cause: unknown) {
+    super(
+      cause instanceof Error
+        ? cause.message
+        : "CATMAID reroot source-state refresh failed.",
+    );
+    this.name = "CatmaidRerootSourceStateRefreshError";
+  }
+}
+
 function getResolvedNodeContextForEdit(
   layer: SegmentationUserLayer,
   stableNodeId: number,
@@ -1737,10 +1774,27 @@ class RerootCommand implements SpatialSkeletonCommand {
     if (resolvedNode.node.parentNodeId === undefined) {
       return;
     }
-    const result = await this.editOperations.commitReroot({
-      node: resolvedNode.node,
-      segmentNodes: resolvedNode.segmentNodes,
-    });
+    let result: CatmaidSpatialSkeletonRerootResult;
+    try {
+      result = await this.editOperations.commitReroot({
+        node: resolvedNode.node,
+        segmentNodes: resolvedNode.segmentNodes,
+      });
+    } catch (error) {
+      if (!(error instanceof CatmaidRerootSourceStateRefreshError)) {
+        throw error;
+      }
+      resolvedNode.skeletonLayer.invalidateSourceCellsForPositions(
+        collectUniqueNodePositions(resolvedNode.segmentNodes),
+      );
+      this.layer.spatialSkeletonState.invalidateCachedSegments([
+        resolvedNode.node.segmentId,
+      ]);
+      this.layer.markSpatialSkeletonNodeDataChanged({
+        invalidateFullSkeletonCache: false,
+      });
+      throw error;
+    }
     this.layer.spatialSkeletonState.rerootCachedSegment(
       resolvedNode.node.nodeId,
     );
@@ -2465,10 +2519,34 @@ export class CatmaidSpatialSkeletonEditCommands {
   private commitReroot(
     request: CatmaidSpatialSkeletonRerootRequest,
   ): Promise<CatmaidSpatialSkeletonRerootResult> {
-    return this.client.rerootSkeleton(
+    return this.commitRerootAndRefreshSourceStates(request);
+  }
+
+  private async commitRerootAndRefreshSourceStates(
+    request: CatmaidSpatialSkeletonRerootRequest,
+  ): Promise<CatmaidSpatialSkeletonRerootResult> {
+    const affectedNodeIds = getSpatiallyIndexedSkeletonPathToRoot(
+      request.segmentNodes,
+      request.node,
+    ).map((node) => node.nodeId);
+    const result = await this.client.rerootSkeleton(
       request.node.nodeId,
       buildCatmaidRerootEditContext(request.node, request.segmentNodes),
     );
+    let nodeSourceStateUpdates: readonly CatmaidSpatialSkeletonNodeSourceStateUpdate[];
+    try {
+      nodeSourceStateUpdates = await getFreshRerootSourceStateUpdates(
+        this.client,
+        request.node.segmentId,
+        affectedNodeIds,
+      );
+    } catch (error) {
+      throw new CatmaidRerootSourceStateRefreshError(error);
+    }
+    return {
+      ...result,
+      nodeSourceStateUpdates,
+    };
   }
 
   private commitDescription(