I think I have a working document cache solution that's actually pretty good.

src/context/document/DocumentContext.tsx

@@ -1,59 +1,28 @@
-import { pb } from "@/lib/pocketbase";
-import { type AnyDocument, type DocumentId } from "@/lib/types";
-import type { RecordModel } from "pocketbase";
 import type { ReactNode } from "react";
-import { createContext, useContext, useEffect, useReducer } from "react";
+import { createContext, useReducer } from "react";
 import type { DocumentAction } from "./actions";
 import { reducer } from "./reducer";
-import { loading, type DocumentState } from "./state";
+import { initialState, type DocumentState } from "./state";
 
-type DocumentContextValue = {
-  state: DocumentState<AnyDocument>;
-  dispatch: (action: DocumentAction<AnyDocument>) => void;
+export type DocumentContextValue = {
+  cache: DocumentState;
+  dispatch: (action: DocumentAction) => void;
 };
 
-const DocumentContext = createContext<DocumentContextValue | undefined>(
+export const DocumentContext = createContext<DocumentContextValue | undefined>(
   undefined,
 );
 
 /**
  * Provider for the record cache context. Provides a singleton RecordCache instance to children.
  */
-export function DocumentProvider({
-  documentId,
-  children,
-}: {
-  documentId: DocumentId;
-  children: ReactNode;
-}) {
-  const [state, dispatch] = useReducer(reducer, loading());
+export function DocumentProvider({ children }: { children: ReactNode }) {
+  const [state, dispatch] = useReducer(reducer, initialState());
 
-  useEffect(() => {
-    async function fetchDocumentAndRelations() {
-      const doc: AnyDocument = await pb
-        .collection("documents")
-        .getOne(documentId, {
-          expand:
-            "relationships_via_primary,relationships_via_primary.secondary",
-        });
-
-      dispatch({
-        type: "ready",
-        doc,
-        relationships: doc.expand?.relationships_via_primary || [],
-        relatedDocuments:
-          doc.expand?.relationships_via_primary?.flatMap(
-            (r: RecordModel) => r.expand?.secondary,
-          ) || [],
-      });
-    }
-
-    fetchDocumentAndRelations();
-  }, [documentId]);
   return (
     <DocumentContext.Provider
       value={{
-        state,
+        cache: state,
         dispatch,
       }}
     >
@@ -61,10 +30,3 @@ export function DocumentProvider({
     </DocumentContext.Provider>
   );
 }
-
-export function useDocument(): DocumentContextValue {
-  const ctx = useContext(DocumentContext);
-  if (!ctx)
-    throw new Error("useDocument must be used within a DocumentProvider");
-  return ctx;
-}

src/context/document/DocumentLoader.tsx

@@ -0,0 +1,48 @@
+import { pb } from "@/lib/pocketbase";
+import { type AnyDocument, type DocumentId } from "@/lib/types";
+import type { RecordModel } from "pocketbase";
+import type { ReactNode } from "react";
+import { useEffect } from "react";
+import { useDocumentCache } from "./hooks";
+
+/**
+ * Provider for the record cache context. Provides a singleton RecordCache instance to children.
+ */
+export function DocumentLoader({
+  documentId,
+  children,
+}: {
+  documentId: DocumentId;
+  children: ReactNode;
+}) {
+  const { dispatch } = useDocumentCache();
+
+  useEffect(() => {
+    async function fetchDocumentAndRelations() {
+      dispatch({
+        type: "loadingDocument",
+        docId: documentId,
+      });
+      const doc: AnyDocument = await pb
+        .collection("documents")
+        .getOne(documentId, {
+          expand:
+            "relationships_via_primary,relationships_via_primary.secondary",
+        });
+
+      dispatch({
+        type: "setDocumentTree",
+        doc,
+        relationships: doc.expand?.relationships_via_primary || [],
+        relatedDocuments:
+          doc.expand?.relationships_via_primary?.flatMap(
+            (r: RecordModel) => r.expand?.secondary,
+          ) || [],
+      });
+    }
+
+    fetchDocumentAndRelations();
+  }, [documentId]);
+
+  return children;
+}

src/context/document/actions.ts

@@ -1,14 +1,9 @@
-import type { AnyDocument, Relationship } from "@/lib/types";
+import type { AnyDocument, DocumentId, Relationship } from "@/lib/types";
 
-export type DocumentAction<D extends AnyDocument> =
+export type DocumentAction =
   | {
-      type: "loading";
-    }
-  | {
-      type: "ready";
-      doc: D;
-      relationships: Relationship[];
-      relatedDocuments: AnyDocument[];
+      type: "loadingDocument";
+      docId: DocumentId;
     }
   | {
       type: "setDocument";
@@ -16,5 +11,12 @@ export type DocumentAction<D extends AnyDocument> =
     }
   | {
       type: "setRelationship";
+      docId: DocumentId;
       relationship: Relationship;
+    }
+  | {
+      type: "setDocumentTree";
+      doc: AnyDocument;
+      relationships: Relationship[];
+      relatedDocuments: AnyDocument[];
     };

src/context/document/hooks.ts

@@ -0,0 +1,23 @@
+import type { DocumentId } from "@/lib/types";
+import { useContext } from "react";
+import { DocumentContext } from "./DocumentContext";
+
+export function useDocument(id: DocumentId) {
+  const ctx = useContext(DocumentContext);
+  if (!ctx)
+    throw new Error("useDocument must be used within a DocumentProvider");
+  return {
+    docResult: ctx.cache.documents[id],
+    dispatch: ctx.dispatch,
+  };
+}
+
+export function useDocumentCache() {
+  const ctx = useContext(DocumentContext);
+  if (!ctx)
+    throw new Error("useDocument must be used within a DocumentProvider");
+  return {
+    cache: ctx.cache,
+    dispatch: ctx.dispatch,
+  };
+}

src/context/document/reducer.ts

@@ -1,71 +1,88 @@
-import _ from "lodash";
-import type {
-  AnyDocument,
-  DocumentId,
-  Relationship,
-  RelationshipType,
-} from "@/lib/types";
+import type { AnyDocument, DocumentId, Relationship } from "@/lib/types";
 import type { DocumentAction } from "./actions";
-import type { DocumentState } from "./state";
+import { ready, loading, unloaded, type DocumentState } from "./state";
+import { relationshipsForDocument } from "@/lib/relationships";
 
-function ifStatus<D extends AnyDocument, S extends DocumentState<D>["status"]>(
-  status: S,
-  state: DocumentState<D>,
-  newState: (state: DocumentState<D> & { status: S }) => DocumentState<D>,
-) {
-  // TODO: Is there a better way to express the type of type narrowing?
-  return state.status === status
-    ? newState(state as DocumentState<D> & { status: S })
-    : state;
+function setLoadingDocument(
+  docId: DocumentId,
+  state: DocumentState,
+): DocumentState {
+  return {
+    ...state,
+    documents: {
+      ...state.documents,
+      [docId]: loading(),
+    },
+  };
 }
 
-export function reducer<D extends AnyDocument>(
-  state: DocumentState<D>,
-  action: DocumentAction<D>,
-): DocumentState<D> {
-  switch (action.type) {
-    case "loading":
-      return {
-        status: "loading",
-      };
-    case "ready":
-      return {
-        status: "ready",
-        doc: action.doc,
-        relationships: _.keyBy(action.relationships, (r) => r.type) as Record<
-          RelationshipType,
-          Relationship
-        >,
-        relatedDocs: _.keyBy(action.relatedDocuments, (r) => r.id) as Record<
-          DocumentId,
-          AnyDocument
-        >,
-      };
+function setDocument(state: DocumentState, doc: AnyDocument): DocumentState {
+  const previous = state.documents[doc.id];
+  const relationships =
+    previous?.type === "ready"
+      ? previous.value.relationships
+      : Object.fromEntries(
+          relationshipsForDocument(doc).map((relationshipType) => [
+            relationshipType,
+            unloaded(),
+          ]),
+        );
 
-    case "setDocument":
-      return ifStatus("ready", state, (state) => {
-        if (state.doc.id === action.doc.id) {
-          return {
-            ...state,
-            doc: action.doc as D,
-          };
-        }
+  return {
+    ...state,
+    documents: {
+      ...state.documents,
+      [doc.id]: ready({
+        doc: doc,
+        relationships,
+      }),
+    },
+  };
+}
 
-        return {
-          ...state,
-          relatedDocs: {
-            ...state.relatedDocs,
-            [action.doc.id]: action.doc,
-          },
-        };
-      });
-    case "setRelationship":
-      return ifStatus("ready", state, (state) => ({
-        ...state,
+function setRelationship(
+  docId: DocumentId,
+  state: DocumentState,
+  relationship: Relationship,
+): DocumentState {
+  const previousResult = state.documents[docId];
+  if (previousResult?.type !== "ready") {
+    return state;
+  }
+  const previousEntry = previousResult.value;
+  return {
+    ...state,
+    documents: {
+      ...state.documents,
+      [docId]: ready({
+        ...previousEntry,
         relationships: {
-          ...state.relationships,
-          [action.relationship.type]: action.relationship,
+          ...previousEntry.relationships,
+          [relationship.type]: ready(relationship),
         },
-      }));
+      }),
+    },
+  };
+}
+
+export function reducer(
+  state: DocumentState,
+  action: DocumentAction,
+): DocumentState {
+  switch (action.type) {
+    case "loadingDocument":
+      return setLoadingDocument(action.docId, state);
+    case "setDocument":
+      return setDocument(state, action.doc);
+    case "setRelationship":
+      return setRelationship(action.docId, state, action.relationship);
+    case "setDocumentTree":
+      return action.relatedDocuments.reduce(
+        setDocument,
+        action.relationships.reduce(
+          setRelationship.bind(null, action.doc.id),
+          setDocument(state, action.doc),
+        ),
+      );
   }
 }

src/context/document/state.ts

@@ -5,17 +5,28 @@ import type {
   RelationshipType,
 } from "@/lib/types";
 
-export type DocumentState<D extends AnyDocument> =
-  | {
-      status: "loading";
-    }
-  | {
-      status: "ready";
-      doc: D;
-      relationships: Record<RelationshipType, Relationship>;
-      relatedDocs: Record<DocumentId, AnyDocument>;
-    };
+export type Result<V> =
+  | { type: "unloaded" }
+  | { type: "error"; err: unknown }
+  | { type: "loading" }
+  | { type: "ready"; value: V };
 
-export const loading = <D extends AnyDocument>(): DocumentState<D> => ({
-  status: "loading",
-});
+export const unloaded = (): Result<any> => ({ type: "unloaded" });
+export const error = (err: unknown): Result<any> => ({ type: "error", err });
+export const loading = (): Result<any> => ({ type: "loading" });
+export const ready = <V>(value: V): Result<V> => ({ type: "ready", value });
+
+export type DocumentState = {
+  documents: Record<
+    DocumentId,
+    Result<{
+      doc: AnyDocument;
+      relationships: Record<RelationshipType, Result<Relationship>>;
+    }>
+  >;
+};
+
+export const initialState = (): DocumentState =>
+  ({
+    documents: {},
+  }) as DocumentState;