Time Travel for Your State: Undo/Redo - with Zustand and React Query (Part 2)

#react #javascript #typescript #webdev

In Part 1, we looked at the classic snapshot approach, storing full copies of state in past / present / future.

That works well in some setups, but in my case, it started to feel heavy.

I’m using Zustand for client-side state and TanStack Query for server state. Since React Query already handles caching and avoids unnecessary refetches, storing full snapshots for every undo step felt redundant.

Instead of storing what the state was, what if we stored how to get back to it?

The Command Pattern (Kind Of)

I came across Rocicorp’s undo library, and it felt like the idea I had in my head but couldn’t quite put into code. And there it was! Simple, elegant and powerful.

Instead of snapshots, you store a pair of functions:

one to undo the action
one to redo it

{
  entries: [
    { undo: fn1, redo: fn2 },
    { undo: fn3, redo: fn4 },
    { undo: fn5, redo: fn6 }
  ],
  index: 1
}

Each entry represents a single action.

undo → how to revert it
redo → how to apply it again

The index tells you where you are in history.

Undo → run entries[index].undo() and move back
Redo → move forward and run entries[index].redo()

Instead of storing data, you’re storing behavior.

Building the UndoManager

Here’s the core class:

type Entry = {
  groupId?: number;
  undo: () => Promise<void>;
  redo: () => Promise<void>;
};

class UndoManager {
  private state = {
    entries: [] as Entry[],
    index: -1,  // -1 means no history yet
    isGrouping: false,
    lastGroupId: 0,
  };

  get canUndo(): boolean {
    return this.state.index >= 0;
  }

  get canRedo(): boolean {
    return this.state.index < this.state.entries.length - 1;
  }

  async add(options: { undo: () => Promise<void>; redo: () => Promise<void> }) {
    // Clear any "future" entries
    this.state.entries.splice(this.state.index + 1);

    // Add new entry
    this.state.entries.push({
      undo: options.undo,
      redo: options.redo,
      groupId: this.state.isGrouping ? this.state.lastGroupId : undefined,
    });

    this.state.index += 1;
  }

  async undo() {
    if (!this.canUndo) return;

    const entry = this.state.entries[this.state.index];
    this.state.index -= 1;
    await entry.undo();
  }

  async redo() {
    if (!this.canRedo) return;

    const entry = this.state.entries[this.state.index + 1];
    this.state.index += 1;
    await entry.redo();
  }
}

Integrating with React Query

To make this usable, I wrapped my mutations so they automatically register undo/redo behavior:

const runWithSnapshotUndo = async (execute) => {
  // Capture state before
  const beforeSnapshot = await getSnapshot();

  // Run the actual mutation
  const result = await execute();

  // Capture state after
  const afterSnapshot = await getSnapshot();

  // Compare - if nothing changed, don't add to history
  if (areSnapshotsEqual(beforeSnapshot, afterSnapshot)) {
    return result;
  }

  // Add undo/redo functions
  await undoManager.add({
    undo: () => applySnapshot(beforeSnapshot),
    redo: () => applySnapshot(afterSnapshot),
  });

  return result;
};

Now any mutation can be wrapped like this:

const addTodoWithUndo = (payload) =>
  runWithSnapshotUndo(() => 
    addTodoMutation.mutateAsync(payload)
  );

And there it is! You’re not storing snapshots as history entries anymore. You’re storing functions—and those functions decide how to restore state.

Bonus 1: Grouping Actions

Sometimes one user action triggers multiple updates.

For example, dragging a task might:

update one item’s position
shift others around it

From the user’s perspective, that’s a single action.

Grouping lets you treat it that way:

undoManager.startGroup();
await updateOperation1();
await updateOperation2();
await updateOperation3();
undoManager.endGroup();

All entries share the same groupId, so one undo reverses the entire sequence.

async undo() {
  if (!this.canUndo) return;

  const entry = this.state.entries[this.state.index];
  this.state.index -= 1;
  await entry.undo();

  const nextEntry = this.state.entries[this.state.index];

  if (
    entry.groupId !== undefined &&
    nextEntry &&
    nextEntry.groupId === entry.groupId
  ) {
    await this.undo();
  }
}

Bonus 2: Keyboard Shortcuts

Since this is a standalone manager, adding keyboard shortcuts is straightforward:

attachKeyboardShortcuts(getContainer: () => HTMLElement | null) {
  const onKey = (e: KeyboardEvent) => {
    const el = getContainer();
    if (!el || !el.contains(e.target)) return;

    const tag = (e.target as HTMLElement)?.tagName;
    if (tag === 'INPUT' || tag === 'TEXTAREA') return;

    if ((e.metaKey || e.ctrlKey) && e.key === 'z' && !e.shiftKey) {
      e.preventDefault();
      void this.undo();
    } else if (
      (e.metaKey || e.ctrlKey) &&
      (e.key === 'y' || (e.key === 'z' && e.shiftKey))
    ) {
      e.preventDefault();
      void this.redo();
    }
  };

  window.addEventListener('keydown', onKey);
  return () => window.removeEventListener('keydown', onKey);
}

What I Like About This

This approach feels a lot more flexible:

You can add undo to specific mutations without touching everything
Async operations work naturally
Grouping is straightforward
Memory usage stays under control

Note: History can still grow large. It’s a good idea to cap it (I usually limit it to around 1,000 entries).

Wrapping Up

Both approaches still have their place:

Snapshot model → great for centralized, predictable state
Command model → better for async logic and mixed state sources

In this setup, the command-style approach ended up fitting better.

Instead of building your app around undo/redo, you layer it on top—only where you need it.

If you’ve implemented undo/redo differently, I’d love to hear how you approached it. I don't have a separate project setup for this, but if you want some reference, feel free to check out my actual project here.

Until next time 👋