Building Custom Middleware for Cross-Cutting Concerns

Hello! Let's dive into the next lesson on Zustand.

In our previous lessons, you've become proficient with Zustand's built-in middleware, using devtools for debugging and persist for saving state. We established that middleware are powerful wrappers that enhance your store's functionality and that their order of application is significant.

Today, we'll move from being a consumer of middleware to a creator. The learning outcome for this lesson is to create a custom middleware to implement cross-cutting concerns like logging state changes. We'll deconstruct the anatomy of a middleware function and build a practical logger from scratch. This skill allows you to extend Zustand to fit any unique requirement of your application, such as custom analytics, error reporting, or specialized data handling.

---

1. The Anatomy of a Zustand Middleware

At its core, a Zustand middleware is a higher-order function. It's a function that takes the store's configuration as an argument and returns a new, enhanced configuration.

The fundamental signature looks like this:

const myMiddleware = (config) => (set, get, api) => {
  // ... middleware logic ...
  return config(
    (args) => { /* new 'set' logic */ }, 
    get, 
    api
  );
};

Let's break this down:

1. myMiddleware is a function that accepts one argument, config. This config is your original store creator function, i.e., (set, get, api) => ({...}).
2. myMiddleware returns another function. Zustand's create will invoke this returned function, passing it the original set, get, and api objects.
3. The magic happens inside. The middleware's main job is to call the original config function. However, it can intercept and modify the arguments passed to it, most commonly the set function.
4. By creating a new version of set, you can execute code before and/or after the actual state update occurs.

This pattern is central to implementing any cross-cutting concern, where you want to apply the same logic across different actions without repeating code in every action.

2. Creating a Custom Logging Middleware

Logging state changes is a classic use case for middleware. It's invaluable for debugging during development, allowing you to see every state transition as it happens.

Let's build a logger middleware that prints the action being applied and the resulting new state to the console.

As you saw in the resource, the implementation is straightforward. Here is the code with a breakdown:

// Our custom logger middleware
const logger = (config) => (set, get, api) =>
  config(
    // This is our new, enhanced 'set' function
    (args) => {
      console.log("  applying", args); // Log what's being set
      set(args); // Call the original 'set' to update the state
      console.log("  new state", get()); // Log the new state using 'get'
    },
    get,
    api
  );

// Now, let's apply it to a store
const useBearStore = create(logger((set) => ({
    bears: 0,
    increase: () => set((state) => ({ bears: state.bears + 1 })),
    removeAllBears: () => set({ bears: 0 }),
})));

How it works:

1. We wrap our store definition (set) => ({...}) with our logger middleware.
2. When increase() is called, it invokes set((state) => ({ bears: state.bears + 1 })).
3. Our middleware intercepts this. The args variable becomes the function (state) => ({ bears: state.bears + 1 }).
4. console.log(" applying", args) prints this function to the console.
5. set(args) executes the original set function, updating the state.
6. console.log(" new state", get()) then fetches the complete, updated state and prints it.

This gives you a clear, real-time trace of your store's activity directly in your browser's console.

You can also design the logger differently. For example, you might prefer to see the state before and after the change. The article "Managing Global State in Web Apps with Zustand React" shows this alternative pattern.

// An alternative implementation from the Dhiwise article
const logMiddleware = (config) => (set, get, api) => config((args) => {
  console.log('State before update:', get());
  set(args);
  console.log('State after update:', get());
});

This flexibility is the key benefit of creating your own middleware—you can tailor it precisely to your needs.

3. Beyond Logging: Middleware for Async Actions

Custom middleware isn't limited to logging. You can implement more complex logic. A common pattern is to create a middleware that helps manage asynchronous actions, similar to how redux-thunk works.

Let's examine a middleware that checks if an action is a function. If it is, the middleware executes it, giving it access to set and get. This allows you to encapsulate async logic cleanly.

Here's the core of the apiMiddleware:

const apiMiddleware = (config) => (set, get, api) =>
  config(
    (action) => {
      if (typeof action === 'function') {
        // If it's a function, execute it.
        // This is where you'd handle an async operation.
        return action(set, get, api);
      }
      // Otherwise, it's a normal state object.
      set(action);
    },
    get,
    api
  );

This pattern allows you to define an action like fetchProducts as an async function directly in your store. When you call fetchProducts(), it's actually calling set(async () => { ... }) under the hood (or a variation of it). The apiMiddleware intercepts this, executes your async function, and gives it the tools (set, get) to update the state when the API call completes.

4. Composing Middleware

Just as you combined devtools and persist, you can chain multiple custom middlewares together. The order of composition determines the flow of execution.

For example, if you wanted to log all actions, including the async thunks handled by apiMiddleware, you would place the logger on the outside.

import { create } from 'zustand';

// Assume 'logger' and 'apiMiddleware' are defined as above

const useProductStore = create(
  logger( // Outer middleware
    apiMiddleware( // Inner middleware
      (set) => ({
        products: [],
        // ... store definition using the async pattern
      })
    )
  )
);

In this setup:

1. An action is dispatched.
2. The logger receives it first, printing " applying", [Function: fetchProducts].
3. It passes the action to apiMiddleware.
4. apiMiddleware identifies it as a function, executes it, and the async logic runs.
5. Any set calls made inside the async function will also be intercepted by the logger, giving you a complete trace.

---

Conclusion

You have now learned how to build your own Zustand middleware, moving beyond the standard library to create tailored solutions for your application's specific needs. This is a powerful technique for keeping your codebase clean and adhering to the Don't Repeat Yourself (DRY) principle.

Key Takeaways:

• A custom middleware is a higher-order function that wraps the store's configuration: (config) => (set, get, api) => config(...).
• The primary mechanism is to provide a new set function that adds behavior before or after calling the original set.
• Logging is a canonical example of a "cross-cutting concern" that is perfectly suited for middleware.
• Middleware can enable advanced patterns, such as a thunk-like system for handling asynchronous actions.
• Custom middleware can be composed with other middleware, and the order of composition is crucial.

Next Lesson Preview:

So far, we've focused on how to define, update, and debug state. In our next lesson, we will shift our attention to a critical aspect of state management: performance. We will analyze how Zustand's selective subscriptions prevent unnecessary component re-renders, one of its key advantages over other libraries like React Context.