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

src/components/DocumentList.tsx

@@ -1,13 +1,12 @@
+import * as Icons from "@/components/Icons.tsx";
 import type { AnyDocument, DocumentId } from "@/lib/types";
 import {
   Dialog,
   DialogPanel,
-  DialogTitle,
   Transition,
   TransitionChild,
 } from "@headlessui/react";
 import { Fragment, useCallback, useState } from "react";
-import * as Icons from "@/components/Icons.tsx";
 
 type Props<T extends AnyDocument> = {
   title: React.ReactNode;
@@ -75,13 +74,13 @@ export function DocumentList<T extends AnyDocument>({
       {error && (
         <div className="bg-red-900 rounded p-4 text-slate-100">{error}</div>
       )}
-      <ul className="space-y-2">
+      <ul className="flex flex-col space-y-2">
         {items.map((item) => (
           <li
             key={item.id}
-            className="bg-slate-800 rounded p-4 text-slate-100 flex flex-row justify-between items-center"
+            className="p-4 bg-slate-800 rounded text-slate-100 flex flex-row justify-between items-center"
           >
-            <div>{renderRow(item)}</div>
+            {renderRow(item)}
 
             {isEditing && (
               <div>

src/components/RelationshipList.tsx

@@ -1,5 +1,5 @@
 import { DocumentList } from "@/components/DocumentList";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache, useDocument } from "@/context/document/hooks";
 import { pb } from "@/lib/pocketbase";
 import { displayName } from "@/lib/relationships";
 import type {
@@ -28,17 +28,30 @@ export function RelationshipList({
 }: RelationshipListProps) {
   const [_loading, setLoading] = useState(true);
   const [error, setError] = useState<string | null>(null);
-  const { state, dispatch } = useDocument();
+  const { docResult, dispatch } = useDocument(root.id);
+  const { cache } = useDocumentCache();
 
-  if (state.status === "loading") {
+  if (docResult.type !== "ready") {
     return <Loader />;
   }
 
+  const doc = docResult.value.doc;
+
   console.info("Rendering relationship list: ", relationshipType);
 
-  const relationship = state.relationships[relationshipType];
-  const itemIds = relationship?.secondary ?? [];
-  const items = itemIds.map((id) => state.relatedDocs[id]).filter((d) => !!d);
+  const relationshipResult = docResult.value.relationships[relationshipType];
+
+  if (relationshipResult?.type !== "ready") {
+    return <Loader />;
+  }
+
+  const relationship = relationshipResult.value;
+
+  const itemIds = relationship.secondary ?? [];
+  const items = itemIds
+    .map((id) => cache.documents[id])
+    .filter((d) => d && d.type === "ready")
+    .map((d) => d.value.doc);
 
   const handleCreate = async (doc: AnyDocument) => {
     setLoading(true);
@@ -54,6 +67,7 @@ export function RelationshipList({
           });
         dispatch({
           type: "setRelationship",
+          docId: doc.id,
           relationship: updatedRelationship,
         });
       } else {
@@ -67,6 +81,7 @@ export function RelationshipList({
           });
         dispatch({
           type: "setRelationship",
+          docId: doc.id,
           relationship: updatedRelationship,
         });
       }
@@ -91,6 +106,7 @@ export function RelationshipList({
           });
         dispatch({
           type: "setRelationship",
+          docId: doc.id,
           relationship: updatedRelationship,
         });
       }

src/components/documents/DocumentView.tsx

@@ -1,24 +1,35 @@
 import { RelationshipList } from "@/components/RelationshipList";
 import { DocumentEditForm } from "@/components/documents/DocumentEditForm";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocument } from "@/context/document/hooks";
 import { displayName, relationshipsForDocument } from "@/lib/relationships";
 import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@headlessui/react";
 import { Link } from "@tanstack/react-router";
 import { Loader } from "../Loader";
+import type { DocumentId } from "@/lib/types";
 
-export function DocumentView() {
-  const { state } = useDocument();
+export function DocumentView({ documentId }: { documentId: DocumentId }) {
+  const { docResult } = useDocument(documentId);
 
-  if (state.status === "loading") {
+  console.info(`Rendering document: `, docResult);
+
+  if (docResult?.type !== "ready") {
     return <Loader />;
   }
 
-  const doc = state.doc;
+  const doc = docResult.value.doc;
 
   const relationshipList = relationshipsForDocument(doc);
 
   return (
     <div key={doc.id} className="max-w-xl mx-auto py-2 px-4">
+      <div>
+        <Link
+          to="/document/$documentId/print"
+          params={{ documentId: doc.id }}
+          className="text-slate-400 hover:text-violet-400 text-sm underline underline-offset-2 transition-colors mb-4"
+        >
+          Back to campaign
+        </Link>
         <Link
           to="/document/$documentId/print"
           params={{ documentId: doc.id }}
@@ -26,6 +37,7 @@ export function DocumentView() {
         >
           Print
         </Link>
+      </div>
       <DocumentEditForm document={doc} />
       <TabGroup>
         <TabList className="flex flex-row flex-wrap gap-1 mt-2">

src/components/documents/location/LocationEditForm.tsx

@@ -1,15 +1,17 @@
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
 import { pb } from "@/lib/pocketbase";
 import type { Location } from "@/lib/types";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders an editable location form
  */
 export const LocationEditForm = ({ location }: { location: Location }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   async function saveLocationName(name: string) {
-    const updated: Location = await pb.collection("documents").update(location.id, {
+    const updated: Location = await pb
+      .collection("documents")
+      .update(location.id, {
         data: {
           ...location.data,
           name,
@@ -19,7 +21,9 @@ export const LocationEditForm = ({ location }: { location: Location }) => {
   }
 
   async function saveLocationDescription(description: string) {
-    const updated: Location = await pb.collection("documents").update(location.id, {
+    const updated: Location = await pb
+      .collection("documents")
+      .update(location.id, {
         data: {
           ...location.data,
           description,

src/components/documents/location/NewLocationForm.tsx

@@ -4,7 +4,7 @@ import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { MultiLineInput } from "@/components/form/MultiLineInput";
 import { SingleLineInput } from "@/components/form/SingleLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new location. Calls onCreate with the new location document.
@@ -16,7 +16,7 @@ export const NewLocationForm = ({
   campaign: CampaignId;
   onCreate: (location: Location) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [name, setName] = useState("");
   const [description, setDescription] = useState("");
   const [adding, setAdding] = useState(false);

src/components/documents/monsters/MonsterEditForm.tsx

@@ -1,15 +1,17 @@
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
 import { pb } from "@/lib/pocketbase";
 import type { Monster } from "@/lib/types";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders an editable monster row
  */
 export const MonsterEditForm = ({ monster }: { monster: Monster }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   async function saveMonsterName(name: string) {
-    const updated = await pb.collection("documents").update(monster.id, {
+    const updated: Monster = await pb
+      .collection("documents")
+      .update(monster.id, {
         data: {
           ...monster.data,
           name,

src/components/documents/monsters/NewMonsterForm.tsx

@@ -3,7 +3,7 @@ import type { CampaignId, Monster } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { SingleLineInput } from "@/components/form/SingleLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new monster. Calls onCreate with the new monster document.
@@ -15,7 +15,7 @@ export const NewMonsterForm = ({
   campaign: CampaignId;
   onCreate: (monster: Monster) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [name, setName] = useState("");
   const [adding, setAdding] = useState(false);
   const [error, setError] = useState<string | null>(null);

src/components/documents/npc/NewNpcForm.tsx

@@ -4,7 +4,7 @@ import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { SingleLineInput } from "@/components/form/SingleLineInput";
 import { MultiLineInput } from "@/components/form/MultiLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new npc. Calls onCreate with the new npc document.
@@ -16,7 +16,7 @@ export const NewNpcForm = ({
   campaign: CampaignId;
   onCreate: (npc: Npc) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [name, setName] = useState("");
   const [description, setDescription] = useState("");
   const [adding, setAdding] = useState(false);

src/components/documents/npc/NpcEditForm.tsx

@@ -1,5 +1,5 @@
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 import { pb } from "@/lib/pocketbase";
 import type { Npc } from "@/lib/types";
 
@@ -7,7 +7,7 @@ import type { Npc } from "@/lib/types";
  * Renders an editable npc form
  */
 export const NpcEditForm = ({ npc }: { npc: Npc }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   async function saveNpcName(name: string) {
     const updated: Npc = await pb.collection("documents").update(npc.id, {
       data: {

src/components/documents/scene/NewSceneForm.tsx

@@ -5,7 +5,7 @@ import type { CampaignId, Scene } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { MultiLineInput } from "@/components/form/MultiLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new scene. Calls onCreate with the new scene document.
@@ -17,7 +17,7 @@ export const NewSceneForm = ({
   campaign: CampaignId;
   onCreate: (scene: Scene) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [text, setText] = useState("");
   const [adding, setAdding] = useState(false);
   const [error, setError] = useState<string | null>(null);

src/components/documents/scene/SceneEditForm.tsx

@@ -1,15 +1,13 @@
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
 import { pb } from "@/lib/pocketbase";
 import type { Scene } from "@/lib/types";
-import { useDocument } from "@/context/document/DocumentContext";
-import { useQueryClient } from "@tanstack/react-query";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders an editable scene form
  */
 export const SceneEditForm = ({ scene }: { scene: Scene }) => {
-  const { dispatch } = useDocument();
-  const queryClient = useQueryClient();
+  const { dispatch } = useDocumentCache();
 
   async function saveScene(text: string) {
     const updated: Scene = await pb.collection("documents").update(scene.id, {

src/components/documents/secret/NewSecretForm.tsx

@@ -5,7 +5,7 @@ import type { CampaignId, Secret } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { SingleLineInput } from "@/components/form/SingleLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new secret. Calls onCreate with the new secret document.
@@ -17,7 +17,7 @@ export const NewSecretForm = ({
   campaign: CampaignId;
   onCreate: (secret: Secret) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [newSecret, setNewSecret] = useState("");
   const [adding, setAdding] = useState(false);
   const [error, setError] = useState<string | null>(null);

src/components/documents/secret/SecretEditForm.tsx

@@ -1,6 +1,6 @@
 // Displays a single secret with discovered checkbox and text.
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 import { pb } from "@/lib/pocketbase";
 import type { Secret } from "@/lib/types";
 import { useState } from "react";
@@ -10,7 +10,7 @@ import { useState } from "react";
  * Handles updating the discovered state and discoveredIn relationship.
  */
 export const SecretEditForm = ({ secret }: { secret: Secret }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [checked, setChecked] = useState(
     !!(secret.data as any)?.secret?.discovered,
   );

src/components/documents/secret/SecretToggleRow.tsx

@@ -1,7 +1,7 @@
 // SecretRow.tsx
 // Displays a single secret with discovered checkbox and text.
-import type { Secret, Session } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
+import type { Secret, Session } from "@/lib/types";
 import { useState } from "react";
 
 /**
@@ -21,6 +21,8 @@ export const SecretToggleRow = ({
   const [loading, setLoading] = useState(false);
 
   async function handleChange(e: React.ChangeEvent<HTMLInputElement>) {
+    e.stopPropagation();
+    e.preventDefault();
     const newChecked = e.target.checked;
     setLoading(true);
     setChecked(newChecked);
@@ -59,7 +61,7 @@ export const SecretToggleRow = ({
   }
 
   return (
-    <div className="flex items-center gap-3">
+    <div className="flex items-center justify-stretch gap-3 w-full">
       <input
         type="checkbox"
         checked={checked}
@@ -68,7 +70,7 @@ export const SecretToggleRow = ({
         aria-label="Discovered"
         disabled={loading}
       />
-      <span>{secret.data.text}</span>
+      {secret.data.text}
     </div>
   );
 };

src/components/documents/session/SessionEditForm.tsx

@@ -1,10 +1,10 @@
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 import { pb } from "@/lib/pocketbase";
 import type { Session } from "@/lib/types";
 
 export const SessionEditForm = ({ session }: { session: Session }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
 
   async function saveStrongStart(strongStart: string) {
     const doc: Session = await pb.collection("documents").update(session.id, {

src/components/documents/treasure/NewTreasureForm.tsx

@@ -5,7 +5,7 @@ import type { CampaignId, Treasure } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
 import { BaseForm } from "@/components/form/BaseForm";
 import { SingleLineInput } from "@/components/form/SingleLineInput";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 
 /**
  * Renders a form to add a new treasure. Calls onCreate with the new treasure document.
@@ -17,7 +17,7 @@ export const NewTreasureForm = ({
   campaign: CampaignId;
   onCreate: (treasure: Treasure) => Promise<void>;
 }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [newTreasure, setNewTreasure] = useState("");
   const [adding, setAdding] = useState(false);
   const [error, setError] = useState<string | null>(null);

src/components/documents/treasure/TreasureEditForm.tsx

@@ -1,6 +1,6 @@
 // Displays a single treasure with discovered checkbox and text.
 import { AutoSaveTextarea } from "@/components/AutoSaveTextarea";
-import { useDocument } from "@/context/document/DocumentContext";
+import { useDocumentCache } from "@/context/document/hooks";
 import { pb } from "@/lib/pocketbase";
 import type { Treasure } from "@/lib/types";
 import { useState } from "react";
@@ -10,7 +10,7 @@ import { useState } from "react";
  * Handles updating the discovered state and discoveredIn relationship.
  */
 export const TreasureEditForm = ({ treasure }: { treasure: Treasure }) => {
-  const { dispatch } = useDocument();
+  const { dispatch } = useDocumentCache();
   const [checked, setChecked] = useState(
     !!(treasure.data as any)?.treasure?.discovered,
   );

src/components/documents/treasure/TreasureToggleRow.tsx

@@ -1,7 +1,8 @@
 // TreasureRow.tsx
 // Displays a single treasure with discovered checkbox and text.
-import type { Treasure, Session } from "@/lib/types";
 import { pb } from "@/lib/pocketbase";
+import type { Session, Treasure } from "@/lib/types";
+import { Link } from "@tanstack/react-router";
 import { useState } from "react";
 
 /**
@@ -68,7 +69,13 @@ export const TreasureToggleRow = ({
         aria-label="Discovered"
         disabled={loading}
       />
-      <span>{treasure.data.text}</span>
+      <Link
+        to="/document/$documentId"
+        params={{ documentId: treasure.id }}
+        className="text-lg !no-underline text-slate-100 hover:underline hover:text-violet-400"
+      >
+        {treasure.data.text}
+      </Link>
     </div>
   );
 };

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;
-}
-
-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
-        >,
-      };
-
-    case "setDocument":
-      return ifStatus("ready", state, (state) => {
-        if (state.doc.id === action.doc.id) {
+function setLoadingDocument(
+  docId: DocumentId,
+  state: DocumentState,
+): DocumentState {
   return {
     ...state,
-            doc: action.doc as D,
-          };
-        }
-
-        return {
-          ...state,
-          relatedDocs: {
-            ...state.relatedDocs,
-            [action.doc.id]: action.doc,
+    documents: {
+      ...state.documents,
+      [docId]: loading(),
     },
   };
-      });
-    case "setRelationship":
-      return ifStatus("ready", state, (state) => ({
+}
+
+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(),
+          ]),
+        );
+
+  return {
     ...state,
+    documents: {
+      ...state.documents,
+      [doc.id]: ready({
+        doc: doc,
+        relationships,
+      }),
+    },
+  };
+}
+
+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 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 loading = <D extends AnyDocument>(): DocumentState<D> => ({
-  status: "loading",
-});
+export const initialState = (): DocumentState =>
+  ({
+    documents: {},
+  }) as DocumentState;

src/routes/__root.tsx

@@ -1,4 +1,5 @@
 import { AuthProvider } from "@/context/auth/AuthContext";
+import { DocumentProvider } from "@/context/document/DocumentContext";
 import { ReactQueryDevtools } from "@tanstack/react-query-devtools";
 import { Outlet, createRootRoute } from "@tanstack/react-router";
 import { TanStackRouterDevtools } from "@tanstack/react-router-devtools";
@@ -7,7 +8,9 @@ export const Route = createRootRoute({
   component: () => (
     <>
       <AuthProvider>
+        <DocumentProvider>
           <Outlet />
+        </DocumentProvider>
       </AuthProvider>
       <TanStackRouterDevtools />
       <ReactQueryDevtools buttonPosition="bottom-right" />

src/routes/_app/_authenticated/document.$documentId.tsx

@@ -1,5 +1,5 @@
 import { DocumentView } from "@/components/documents/DocumentView";
-import { DocumentProvider } from "@/context/document/DocumentContext";
+import { DocumentLoader } from "@/context/document/DocumentLoader";
 import type { DocumentId } from "@/lib/types";
 import { createFileRoute } from "@tanstack/react-router";
 
@@ -11,11 +11,10 @@ export const Route = createFileRoute(
 
 function RouteComponent() {
   const { documentId } = Route.useParams();
-  console.info("Rendering document route: ", documentId);
 
   return (
-    <DocumentProvider documentId={documentId as DocumentId}>
-      <DocumentView />
-    </DocumentProvider>
+    <DocumentLoader documentId={documentId as DocumentId}>
+      <DocumentView documentId={documentId as DocumentId} />
+    </DocumentLoader>
   );
 }