DEV Community: Grzegorz Dubiel

The Role of Graph Structure in LLM-Powered Applications

Grzegorz Dubiel — Wed, 24 Sep 2025 09:34:18 +0000

Creating features that rely on LLMs often requires handling complex tasks, since the output of LLMs is not deterministic. This creates numerous edge cases that need to be covered. Sometimes we need to preprocess text or files, validate the output, involve the user, and perform many other actions. Without proper planning and robust, scalable abstractions, our application's logic can quickly turn into chaos. That’s why it’s often best to look for battle-tested patterns as a foundation for AI integration code. I raised a similar topic in my previous article about summarizing large documents with LLMs using the Map-Reduce pattern
. In this article, we’ll take the Map-Reduce flow from that post and organize it using a graph structure to make our code more explicit. Both that post and this one are inspired by LangChain’s tutorial on map-reduce for summarization.

Why use a graph in AI integrations?

As I mentioned, even a simple flow can sometimes become difficult to manage. For example, we might need to call an LLM recursively until we achieve the desired result, or we may have two or more LLMs collaborating on a single task. To standardize these processes and make them more readable and maintainable, we can define functions as nodes in a graph, connected by edges that explicitly specify what should happen after each function call. This approach also allows us to share state between function calls, which is very important, since flows involving LLMs usually require storing the task context.

But let’s start this story from the beginning.

What is a Graph Structure in a Few Words?

A graph is a non-linear data structure made up of nodes connected by edges, which indicate relationships between the nodes. A graph can be undirected or directed, and it can also be circular or acyclic. In general, a graph is a very flexible data structure, commonly used to represent relationships between various entities. This flexibility makes it a perfect fit for our use case.

Imagine we have a function represented by a node. The function calls an LLM, and the LLM returns output that goes to an evaluation node (connected by an edge to the first node). The evaluation node calls a specialized SLM (Small Language Model). If the evaluation passes, the output is sent to the output node; if not, another edge leads back to the first node for a retry, this time with annotations to enhance the initial prompt.

Let’s Get Our Hands Dirty With Code

In the previous article, we implemented the Map-Reduce pattern to summarize a blog by summarizing its scraped articles. Each article was sent to the LLM in parallel for summarization. If the resulting list was too long, it was broken down into sublists and summarized again with a reduce prompt, recursively, to achieve the desired compression. Finally, the collapsed list was summarized once more with a reduce prompt to produce the final condensed summary.

The task is quite complex to organize. We have parallel calls, conditional calls, recursion, and splitting involved. Of course, the solution looks neat, but imagine trying to scale it or explain it to new teammates. That’s why we sometimes need well-known patterns and frameworks to make our code more familiar to other developers.

We will be using a framework created with graphs in mind—langgraph. langgraph provides all the primitives needed to turn our flow into a graph. In my opinion, langgraph is one of the best tools for developing AI integrations at scale. Building on a graph structure is a great idea and makes the flow easier to understand.

Preparing Utils

We need to define some utilities for processing data and running a model:

import { ChatOpenAI } from "@langchain/openai";
import { TokenTextSplitter } from "@langchain/textsplitters";

const MAX_TOKEN = 1000;

const textSplitter = new TokenTextSplitter({
  chunkSize: MAX_TOKEN,
  chunkOverlap: 0,
});
const llm = new ChatOpenAI({
  model: "gpt-5-mini",
  apiKey: process.env.OPENAI_API_KEY,
});

Defining a State

The first major improvement to our code is having a tool for defining a shared state, and even making it type-safe.

The state definition looks like this:

/* rest of the code */
import { Annotation } from "@langchain/langgraph";

/* rest of the code */

const OverallState = Annotation.Root({
  contents: Annotation<string[]>,
  summaries: Annotation<string[]>({
    reducer: (state, update) => state.concat(update),
  }),
  collapsedSummaries: Annotation<Document[]>,
  finalSummary: Annotation<string>,
});

As you can see, we are using a single primitive to define the schema of our state as well as the TypeScript types via a generic. The Annotation.Root method is a wrapper used to define the top-level state structure. The properties of the object supplied to Root are channels that store the data returned from the nodes (functions in our graph). The reducer property is a function used to merge the most recent result returned from a node with the current state. Later on, we will pass this OverallState to the graph builder.

Defining Node Functions

The functions we are going to define reflect the four steps we explored in the previous article, which are:

Pre-processing phase

Data returned from the scraper is passed to our graph invocation.

import { runSitemapBasedScraper } from "./scraper/main";
import graph from "./summarizer/with-langgraph/graph";

async function main() {
  const scrappingResults = await runSitemapBasedScraper([
    "https://www.aboutjs.dev",
  ]);

  const filteredScrapedResults = scrappingResults.filter((result) => {
    if (result.error) {
      console.error(`❌ ${result.url}: ${result.error}`);
    }
    return result.success;
  });

  const summarized = await graph.invoke({
    contents: filteredScrapedResults.flatMap((result) =>
      result.posts.map((post) => post.content),
    ),
  });
  console.log("summarized", summarized.finalSummary);
}

void main();

Don’t worry, we’re not skipping anything. We’ll go through the process of building the graph step by step. In this snippet, I just wanted to show how the input data is passed. Our preprocess node, which will be part of the graph, looks like this:

import { Document } from "@langchain/core/documents";
import { Annotation } from "@langchain/langgraph";

/* rest of the code */

const OverallState = Annotation.Root({
  contents: Annotation<string[]>,
  summaries: Annotation<string[]>({
    reducer: (state, update) => state.concat(update),
  }),
  collapsedSummaries: Annotation<Document[]>,
  finalSummary: Annotation<string>,
});

/* rest of the code */

const preProcess = async (state: typeof OverallState.State) => {
  const docs = state.contents.map(
    (content) => new Document({ pageContent: content }),
  );
  const splitDocs = await textSplitter.splitDocuments(docs);
  return { contents: splitDocs.map((doc) => ({ content: doc.pageContent })) };
};
/* rest of the code */

As you can se our node function has a state parameter which is nothing more than our OverallState defined earlier. State is typed. Preprocess function just gets content(posts) passed during invocation and splits into smaller sub-documents to make them to not exceed a given token limit.

Mapping Phase In the mapping phase, we need to generate a summary for each article or its chunk:

Previous functions' implementation

async function runMappers(formattedDocs: Document[]): Promise<string[]> {
  console.log("Summarization started...");
  const splitDocs = await textSplitter.splitDocuments(formattedDocs);

  const results = await model.batch(
    splitDocs.map((doc) => [
      {
        role: "user",
        content: mapTemplate(doc.pageContent),
      },
    ]),
  );

  return results.map((result) => result.content as string);
}

Function's Implementation with langgraph

import { mapPrompt, reducePrompt } from "./prompts";
import { Send } from "@langchain/langgraph";

type SummaryState = {
  content: string;
};

/* REST OF THE CODE */

const mapContents = (state: typeof OverallState.State) => {
  return state.contents.map(
    (content) => new Send("generateSummary", { content }),
  );
};
const generateSummary = async (
  state: SummaryState,
): Promise<{ summaries: string[] }> => {
  const prompt = await mapPrompt.invoke({
    context: state.content,
  });

  const response = await llm.invoke(prompt);
  return { summaries: [String(response.content)] };
};

/* REST OF THE CODE */

The mapContents function from the refactor maps content into a callback that returns a Send object containing the content and the name of the next node (function). Each Send instance spawns one node and transports data from the previous node to that spawned node. As a result, we receive one node per content item. So, mapContents acts like an edge → node generator.

The generateSummary function sends a request to the LLM along with the map prompt. One call is performed per Send. The response from each function is placed in the state channel: summaries. The returned value is added to the array, and the results from each generateSummary call are merged by the reduce callback defined inside Annotation.Root.

Reducing Phase

In this phase, the list of summaries is reduced to the final summary. We also need to handle the process of collapsing the list of summaries. The entire list cannot exceed a given limit, which can be determined by the context window of the LLM or other factors. This is also where we encounter the limitations of the Map-Reduce pattern. Because we need to break the list into sublists, some important pieces of information—expressed, for example, by long paragraphs—can also be split. Therefore, it is important to test how large each chunk should be. In our case, we split the content into very small parts for explanatory purposes.

Let’s take a look at the previous version of the function for the reducing phase:

async function lengthFunction(summaries: string[]) {
  const tokenCounts = await Promise.all(
    summaries.map(async (summary) => {
      return model.getNumTokens(summary);
    }),
  );
  return tokenCounts.reduce((sum, count) => sum + count, 0);
}

export async function splitSummariesByTokenLimit(
  summaries: string[],
  tokenLimit: number,
): Promise<string[][]> {
  const listOfSummariesSublists: string[][] = [];
  let sublist: string[] = [];
  for (const summary of summaries) {
    const chunks = await recursiveTextSplitter.splitText(summary);
    for (const chunk of chunks) {
      const candidateList = [...sublist, chunk];
      const candidateTokens = await lengthFunction(candidateList);
      if (candidateTokens > tokenLimit) {
        if (sublist.length > 0) {
          listOfSummariesSublists.push(sublist);
          sublist = [];
        }
      }
      sublist.push(chunk);
    }
  }
  if (sublist.length > 0) {
    listOfSummariesSublists.push(sublist);
  }
  return listOfSummariesSublists;
}

async function reduceSummariesBatch(listOfSummaries: string[][]) {
  const result = await model.batch(
    listOfSummaries.map((summaries) => [
      {
        role: "user",
        content: reduceTemplate(summaries.join("\n\n")),
      },
    ]),
  );
  return result.map((res) => res.content as string);
}

async function checkShouldCollapse(summaries: string[]) {
  const tokenCount = await lengthFunction(summaries);
  return tokenCount > 1000;
}

async function collapseSummaries(
  summaries: string[],
  recursionLimit = 5,
  iteration = 0,
) {
  console.log("Collapsing summaries...");
  if (summaries.length === 0) {
    return [];
  }
  const splitDocLists = await splitSummariesByTokenLimit(summaries, CHUNK_SIZE);

  const results = await reduceSummariesBatch(splitDocLists);

  let shouldCollapse = await checkShouldCollapse(results);
  if (shouldCollapse && iteration < recursionLimit) {
    console.log("Token count exceeds limit, collapsing summaries further...");
    return collapseSummaries(results, recursionLimit, iteration + 1);
  }
  return results;
}

… and then at the current version of the function for the mapping phase, implemented using langgraph’s primitives:

import {
  collapseDocs,
  splitListOfDocs,
} from "langchain/chains/combine_documents/reduce";
import { reducePrompt } from "./prompts";

/* rest of the code */

const collectSummaries = async (state: typeof OverallState.State) => {
  return {
    collapsedSummaries: state.summaries.map(
      (summary) => new Document({ pageContent: summary }),
    ),
  };
};

async function shouldCollapse(state: typeof OverallState.State) {
  let numTokens = await lengthFunction(state.collapsedSummaries);
  if (numTokens > MAX_TOKEN) {
    return "collapseSummaries";
  } else {
    return "generateFinalSummary";
  }
}
async function reduceSummaries(input: Document[]) {
  const prompt = await reducePrompt.invoke({ docs: input });
  const response = await llm.invoke(prompt);
  return String(response.content);
}

const collapseSummaries = async (state: typeof OverallState.State) => {
  const docLists = splitListOfDocs(
    state.collapsedSummaries,
    lengthFunction,
    MAX_TOKEN,
  );
  const results = [];
  for (const docList of docLists) {
    results.push(await collapseDocs(docList, reduceSummaries));
  }

  return { collapsedSummaries: results };
};

const generateFinalSummary = async (state: typeof OverallState.State) => {
  const response = await reduceSummaries(state.collapsedSummaries);
  return { finalSummary: response };
};

/* rest of the code */

The collectSummaries function takes summaries from the mapping phase (the summaries channel) and copies them to the collapsedSummaries channel, where they can be modified by subsequent functions. In the previous version, this was not necessary.

Next, we have a function for checking whether the list should be collapsed into sublists (checkShouldCollapse), which returns a string indicating whether we should collapse the list or go straight to the final summary. In the previous version, this function returned a boolean, whereas the current version returns a string that indicates which function should be called next, depending on whether the list is too long.

Then, the reduceSummaries function is used to ask an LLM to generate a summary of a list when used in generateFinalSummary, or of a sublist when used inside the collapseSummaries function. In the previous version, we had two functions for sending batch requests to the LLM in the case of collapsing, and just one request when generating the final summary.

In the current version, the process of splitting lists into sublists and reducing them is simplified, as we use utilities from langchain.

Putting It All Together Into a Graph

Alright, this will be the most important part of this refactor, as we will be able to see the benefits of organizing our flow into a graph structure.

First, let’s take a look at the previous version:

/* rest of the code */
export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {
  const formattedDocs = documents.map(
    (doc) =>
      new Document({
        pageContent: doc.content,
        metadata: {
          title: doc.title,
          link: doc.link,
          date: doc.date,
          source: doc.source,
          selector: doc.selector,
          index: doc.index,
        },
      }),
  );

  let summaries = await runMappers(formattedDocs);

  const shouldCollapse = await checkShouldCollapse(summaries);
  if (shouldCollapse) {
    summaries = await collapseSummaries(summaries, maxIterations);
  }
  const finalSummary = await reduceSummaries(summaries);
  console.log("finalSummary", finalSummary);
}
/* rest of the code */

At first glance, everything seems clear and explicit—and it really is. The first time, we might overlook that the collapseSummaries function contains recursion, but we probably notice it on the second iteration. Now imagine that we decide to add more tweaks like this: the main function would be full of implicit actions and could quickly become spaghetti code.

If we don’t plan to expand this logic, it’s fine. But if we have in mind more actions that need to be added to this flow, the graph will really shine:

import { StateGraph } from "@langchain/langgraph";

/* rest of the code */
const graphBuilder = new StateGraph(OverallState)
  // registering nodes
  .addNode("generateSummary", generateSummary)
  .addNode("collectSummaries", collectSummaries)
  .addNode("collapseSummaries", collapseSummaries)
  .addNode("generateFinalSummary", generateFinalSummary)
  .addNode("preProcess", preProcess)
  // preprocessing phase
  .addEdge("__start__", "preProcess")
  // mapping phase
  .addConditionalEdges("preProcess", mapContents, ["generateSummary"])
  .addEdge("generateSummary", "collectSummaries")
  // reducing phase
  .addConditionalEdges("collectSummaries", shouldCollapse, [
    "collapseSummaries",
    "generateFinalSummary",
  ])
  .addConditionalEdges("collapseSummaries", shouldCollapse, [
    "collapseSummaries",
    "generateFinalSummary",
  ])
  .addEdge("generateFinalSummary", "__end__");
const graph = graphBuilder.compile();

export default graph;

Here is the smart and flexible abstraction for our flow.

Firstly, we initialize a StateGraph instance, supplying the state that will be available in all nodes and edges. Then, we use the addNode method to register our core functions responsible for mapping and reducing summaries, passing the node’s name as a string. The node’s name serves as an alias for the function, which is passed as the second argument.

Then we register edges using the addEdge and addConditionalEdges methods. Standard edges act as connections between two functions. For example, in this case: .addEdge("generateSummary", "collectSummaries"), we know that after calling the generateSummary function, the collectSummaries function will be called.

We can also use the addConditionalEdge method to determine which node or nodes should be called. The first argument, as in the standard addEdge method, is the entry node. Then we pass a function that decides which node should come next, and the third argument is the array of possible choices that the conditional function can return. This method also allows us to spawn multiple nodes that can be executed in parallel, and that’s exactly what we do. In the fragment: .addConditionalEdges("preProcess", mapContents, ["generateSummary"]), we go from the preProcess function to the mapContents function, which spawns multiple generateSummary nodes using the Send primitive:

const mapContents = (state: typeof OverallState.State) => {
  return state.contents.map(
    (content) => new Send("generateSummary", { content }),
  );
};

The nodes __start__and __end__are predefined and simply indicate the start and end of the graph.

And now, the last thing: have you spotted the recursion case in our graph? Of course, it appears in this fragment:

/* rest of the code */
  .addConditionalEdges("collectSummaries", shouldCollapse, [
    "collapseSummaries",
    "generateFinalSummary",
  ])
  .addConditionalEdges("collapseSummaries", shouldCollapse, [
    "collapseSummaries",
    "generateFinalSummary",
  ])
  .addEdge("generateFinalSummary", "__end__");
/* rest of the code */

In collapseSummaries, we grab all the summaries from the mapping phase and then check whether they should be collapsed. If not, we move on to generateFinalSummary, which is connected to the __end__ node that marks the end of execution. Otherwise, we jump back to collapseSummaries. Having a structure like that, we can quickly identify a loop just by analyzing the tree.

Thanks to this representation, we can even draw a diagram to visualize the flow in a graph structure:

Now we see that the most benefits we get from organizing a function calling in the graph not refactoring each functions itself. Thanks to this we achieved readability, extendability and the code became explicit.

Summary

With langgraph at hand, we can implement the graph seamlessly. Langgraph provides us with convenient primitives to utilize the graph structure in our flow without writing any low-level code for the graph abstraction. By building frameworks around well-known data structures, we can be confident that every new team member can become productive sooner. In the case of building LLM-powered applications, this is crucial, as we often need to implement complex flows to handle multiple tasks with non-deterministic outputs. Of course, in cases where the process is simple and we know it won’t be expanded in the future, creating additional abstraction is premature, so careful planning is required.

Thanks for reading, and I hope you learned something from my content. Stay tuned for the next article. 👋

PS: Checkout the GitHub repo related to this article.

Turning Entire Blogs into Short Summaries: Map-Reduce for LLMs

Grzegorz Dubiel — Wed, 03 Sep 2025 19:53:36 +0000

Our need for scale is constantly growing. Humanity processes enormous amounts of data. Every day, engineers have to cope with limitations. This reality often forces us to cut corners, creating workarounds to make big things happen. One of the best examples is using LLMs for processing large documents. On one side, we have still-limited technology such as LLMs with restricted context windows, and on the other, we face huge knowledge bases represented by charts, documents, audio, movies, and codebases. In a short time, many useful applications have been created, such as NotebookLM for research work, Claude Code or Cursor for agentic coding. Each of these tools needs to process information that exceeds an LLM’s context window. To make this possible, we must utilize multiple patterns. Let’s consider one of these patterns, called Map-Reduce. To make it easier to understand, we will build a small app for summarizing blogs.

Map-Reduce Pattern

This method was proposed in 2004, by two people from Google: Jeffrey Dean and Sanjay Ghemawat, to process large amounts of distributed data.
In simple terms, it allows a program to efficiently process massive amounts of data in chunks and then reduce it to a final outcome, ensuring speed and overall efficiency of the process.

As the name suggests, the pattern is a combination of two functions:

Map, which concurrently calls a non-void function on each chunk of data (for example, counting expected keywords in the chunk and returning an intermediate key-value pair representing each keyword and its number of occurrences).
Reduce, which takes the output of the map functions and aggregates it into the final result (for example, returning a table with the total number of occurrences for each keyword in the dataset).

Because each map task can run independently on different workers, this process scales easily, and the degree of parallelism can be controlled to balance speed and resource usage.

In the context of using LLMs, this pattern is particularly useful because we always deal with a limited context window. In real-world scenarios, the content or data we want to process often exceeds this window. We cannot process an entire blog with just one or two calls to the LLM. Even if we could, the task would be extremely slow and potentially inaccurate, as LLMs generally reason better on smaller chunks of data. This article and an idea of the app we will build was inspired by the map-reduce pattern tutorial from LangChain’s documentation website.

Project Overview

The first thing I’d like to mention is that in this article, we will try to implement this pattern with minimal help from AI frameworks for building AI solutions like langgraph. The core logic will be implemented by ourselves; we will use only utility functions from langchain to split text and format data.

Our logic will be divided into the following steps:

Pre-processing phase– Collecting selected blog articles using a vibe-coded web scraper.
Mapping phase– Looping over each article, splitting it, and summarizing each part with an LLM in parallel using the map prompt.
Reducing phase– If the summaries produced by the map phase exceed the limit, they are collapsed by splitting them into sub-lists and summarizing each sub-list again with the reduce prompt into a consolidated summary. This reduce phase runs recursively until the mapped or collapsed summaries fit within the limit and can be summarized into the final summary, or until the recursion limit is reached. Otherwise, if the summaries already fit, they are summarized directly with the reduce prompt into the final summary.

Pre-Processing Phase

As I mentioned, to get articles for summarization we will use a scraper. We won’t dive deep into the scraper’s code. We will call the scraper in the main function to fetch all the posts from my blog’s website:

import { summarizeDocuments } from "./summarizer/main";
import { runDirectScraper, runSitemapBasedScraper } from "./scraper/main";

async function main() {
  const scrappingResults = await runSitemapBasedScraper([
    "https://www.aboutjs.dev",
  ]);

  const filteredScrappedResults = scappingResults.filter((result) => {
    if (result.error) {
      console.error(`❌ ${result.url}: ${result.error}`);
    }
    return result.success;
  });
}

void main();

Next, we need to declare a main function for the summarizer — the place where the core logic will run:

// summarizer/types.ts

export type Document = {
  title: string;
  content: string;
  link: string;
  date: string;
  source: string;
  selector: string;
  index: number;
};

// summarizer/main.ts
import type { Document as LocalDocument } from "./type";

export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {}

The function has two parameters: documents, which serve as the input data (articles in our case), and an optional maxIterations parameter that defines how many recursive operations can be performed when the content does not fit into the limit.

Map Phase

The first thing we do is map over all blog posts and prepare them for the text splitter by formatting them into the Document type:

import type { Document as LocalDocument } from "./type";
import { Document } from "@langchain/core/documents";

export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {
  const formattedDocs = documents.map(
    (doc) =>
      new Document({
        pageContent: doc.content,
        metadata: {
          title: doc.title,
          link: doc.link,
          date: doc.date,
          source: doc.source,
          selector: doc.selector,
          index: doc.index,
        },
      }),
  );
}

Now, we prepare a function that will split our documents into chunks and summarize each chunk with the LLM.

To craft it, we first need to prepare a few things:

LLM model runnable

We will use gpt-5-mini. For our summarizer, this model is lightweight and fast — more than good enough for the task. To handle the model instance, we will use the openai integration from langchain.

   /* REST OF THE CODE */

   import { ChatOpenAI } from "@langchain/openai";

   const model = new ChatOpenAI({
     model: "gpt-5-mini",
     apiKey: process.env.OPENAI_API_KEY,
   });

   /* REST OF THE CODE */

prompt

In our case the crucial task of the map phase is to summarize each blog post. We need a prompt that will be sent along with the blog post text, or its chunk, to the LLM for summarization:

   export const mapTemplate = (content: string) => `
   You are an expert content analyzer. Your task is to extract and summarize the key information from the following document.

   Please analyze the content and provide:
   1. Main topics and themes
   2. Key insights and takeaways
   3. Important facts, statistics, or examples
   4. Core concepts or ideas presented

   Format your summary in the bullet points format.

   Summary should be brief and to the point.

   Document Content: ${content}

   Provide a concise but comprehensive summary that captures the essential information from this document. Focus on the most valuable and actionable content.
   `;

The prompt is returned by a template function that takes the blog post content, or its chunk, as an argument. This allows us to dynamically create a prompt with the embedded content for summarization.

text splitter

We need to prepare a tokenTextSplitter, which is capable of splitting a document into smaller sub-documents:

   // summarizer/const.ts
   export const CHUNK_SIZE = 1000;

   // summarizer/main.ts
   import {
     TokenTextSplitter,
     RecursiveCharacterTextSplitter,
   } from "@langchain/textsplitters";
   import { CHUNK_SIZE } from "./const";

   const textSplitter = new TokenTextSplitter({
     chunkSize: CHUNK_SIZE,
     chunkOverlap: 0,
   });

   /* REST OF THE CODE */

We use an extremely small chunk size to illustrate the recursive process of analyzing texts.

Now we are able to put together the function for running the mappers:

import type { Document as LocalDocument } from "./type";
import { TokenTextSplitter } from "@langchain/textsplitters";
import { ChatOpenAI } from "@langchain/openai";
import { mapTemplate } from "./prompts";
import { CHUNK_SIZE } from "./const";
import { Document } from "@langchain/core/documents";

const model = new ChatOpenAI({
  model: "gpt-5-mini",
  apiKey: process.env.OPENAI_API_KEY,
});

const textSplitter = new TokenTextSplitter({
  chunkSize: CHUNK_SIZE,
  chunkOverlap: 0,
});

async function runMappers(formattedDocs: Document[]): Promise<string[]> {
  console.log("Summarization started...");
  const splitDocs = await textSplitter.splitDocuments(formattedDocs);

  const results = await model.batch(
    splitDocs.map((doc) => [
      {
        role: "user",
        content: mapTemplate(doc.pageContent),
      },
    ]),
  );

  return results.map((result) => result.content as string);
}

export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {
  const formattedDocs = documents.map(
    (doc) =>
      new Document({
        pageContent: doc.content,
        metadata: {
          title: doc.title,
          link: doc.link,
          date: doc.date,
          source: doc.source,
          selector: doc.selector,
          index: doc.index,
        },
      }),
  );

  let summaries = await runMappers(formattedDocs);
}

Articles will be split into smaller chunks (sub-documents), and each chunk will then be sent along with the map prompt to the LLM for summarization. The batch method allows us to efficiently send multiple tasks to the LLM in parallel. We can control the maximum number of concurrent requests sent per batch in the LLM model runnable configuration. The results from all the mappers will be returned for processing in the next phase.

Reducing Phase

In the next step, we will take all summaries returned from the map phase and consolidate them into a single final summary. In this step, we also need to handle cases where the summaries from the mappers exceed the LLM’s context window or any other limits. To do this, we will collapse the list of summaries by splitting it into sublists and then summarizing each one. As a result, we will get one summary per sublist.

The first function we create for the reduce phase will determine whether we even need to perform collapsing:

/* REST OF CODE */

async function lengthFunction(summaries: string[]) {
  const tokenCounts = await Promise.all(
    summaries.map(async (summary) => {
      return model.getNumTokens(summary);
    }),
  );
  return tokenCounts.reduce((sum, count) => sum + count, 0);
}

async function checkShouldCollapse(summaries: string[]) {
  const tokenCount = await lengthFunction(summaries);
  return tokenCount > 2000;
}

/* REST OF CODE */

The lengthFunction takes all the summaries and sums up their token count. The checking function then compares this total number of tokens with a hardcoded limit, ensuring that the summaries do not exceed that limit.

Collapsing Summaries

Next, we will build a function for the recursive collapsing of summaries. Collapsing will be done by calling an LLM and requesting a consolidation(summary) of the list of summaries.

Let’s create the reduce prompt for this task:

/* REST OF THE CODE */
export const reduceTemplate = (summaries: string) => `

The following is a set of summaries:
${summaries}
Take these and create one summary as a whole context gathered from the summaries.

Keep it concise and focused on the main points, avoiding unnecessary details. The goal is to distill the essence of the summaries into a single, coherent summary.
`;

Next, we need to prepare two additional functions for splitting the summary list into sublists.

The first will be another token splitter, responsible for handling the edge case where a single summary from the list exceeds the maximum sublist size.

/* REST OF THE CODE */

import {
  TokenTextSplitter,
  RecursiveCharacterTextSplitter,
} from "@langchain/textsplitters";

/* REST OF THE CODE */
const recursiveTextSplitter = new RecursiveCharacterTextSplitter({
  chunkSize: CHUNK_SIZE,
  lengthFunction: (text) => {
    return model.getNumTokens(text);
  },
  chunkOverlap: 0,
});

/* REST OF THE CODE */

The RecursiveCharacterTextSplitter is sufficient in this case, since we focus on a single oversized summary. This splitter recursively breaks the text down from sentences to words, trying to keep fragments as long as possible.

Now we can define a function for splitting the list:

/* REST OF THE CODE */

export async function splitSummariesByTokenLimit(
  summaries: string[],
  tokenLimit: number,
): Promise<string[][]> {
  const listOfSummariesSublists: string[][] = [];
  let sublist: string[] = [];
  for (const summary of summaries) {
    const chunks = await recursiveTextSplitter.splitText(summary);

    for (const chunk of chunks) {
      const candidateList = [...sublist, chunk];
      const candidateTokens = await lengthFunction(candidateList);
      if (candidateTokens > tokenLimit) {
        if (sublist.length > 0) {
          listOfSummariesSublists.push(sublist);
          sublist = [];
        }
      }
      sublist.push(chunk);
    }
  }
  if (sublist.length > 0) {
    listOfSummariesSublists.push(sublist);
  }
  return listOfSummariesSublists;
}

/* REST OF THE CODE */

The function looks quite complicated because the task is not trivial. It iterates over each summary, automatically splitting a summary into smaller chunks if it exceeds the sublist limit. Each potential chunk is then added to a candidate list. The candidate list’s token count is calculated using the lengthFunction defined earlier. If the length exceeds the limit, the current sublist is pushed to the final results array, and a new sublist is started. Otherwise, the chunk is simply added to the current sublist.

We need one more function to collapse the summaries:

/* REST OF THE CODE */
import { mapTemplate, reduceTemplate } from "./prompts";

/* REST OF THE CODE */

async function reduceSummariesBatch(listOfSummaries: string[][]) {
  const result = await model.batch(
    listOfSummaries.map((summaries) => [
      {
        role: "user",
        content: reduceTemplate(summaries.join("\n\n")),
      },
    ]),
  );
  return result.map((res) => res.content as string);
}

/* REST OF THE CODE */

The function runs parallel calls to the LLM to summarize each sublist of summaries. As a result, the list of summaries becomes shorter.

Alright, we can now create the main function for recursively collapsing the summaries:

/* REST OF THE CODE */

async function collapseSummaries(
  summaries: string[],
  recursionLimit = 5,
  iteration = 0,
) {
  console.log("Collapsing summaries...");
  if (summaries.length === 0) {
    return [];
  }
  const splitDocLists = await splitSummariesByTokenLimit(summaries, CHUNK_SIZE);

  const results = await reduceSummariesBatch(splitDocLists);

  let shouldCollapse = await checkShouldCollapse(results);
  if (shouldCollapse && iteration < recursionLimit) {
    console.log("Token count exceeds limit, collapsing summaries further...");
    return collapseSummaries(results, recursionLimit, iteration + 1);
  }
  return results;
}
/* REST OF THE CODE */

In this function, we can easily and clearly organize each step of collapsing, as our complex logic is neatly packed into descriptive functions. First, we split the list of summaries into sublists. Then, we reduce each sublist in a batch using the LLM. Next, we check if the reduced list now fits within the context window. If it does not, we perform another collapse; otherwise, or when the recursion limit is reached, we simply return the reduced list.

Now we can call the collapseSummaries function conditionally inside the main summarizer's function:

/* REST OF THE CODE */

export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {
  const formattedDocs = documents.map(
    (doc) =>
      new Document({
        pageContent: doc.content,
        metadata: {
          title: doc.title,
          link: doc.link,
          date: doc.date,
          source: doc.source,
          selector: doc.selector,
          index: doc.index,
        },
      }),
  );

  let summaries = await runMappers(formattedDocs);

  const shouldCollapse = await checkShouldCollapse(summaries);
  if (shouldCollapse) {
    summaries = await collapseSummaries(summaries, maxIterations);
  }
}
/* REST OF THE CODE */

Final Reduce

After collapsing and ensuring that the list of mapped summaries does not exceed the limit, we can reduce those summaries into a single consolidated final summary.

To do this, we use the reduceTemplate prompt that we used earlier when collapsing sublists of summaries.

Let’s define a function for reducing a single list of summaries:

/* REST OF THE CODE */
async function reduceSummaries(summaries: string[]) {
  const result = await model.invoke([
    {
      role: "user",
      content: reduceTemplate(summaries.join("\n\n")),
    },
  ]);
  return result.content as string;
}
/* REST OF THE CODE */

We did something similar earlier when collapsing each sublist of summaries in batches. Now, we have only one list at a time to reduce.

The last thing we should do is call this function in the main summarizer function:

/* REST OF THE CODE */

export async function summarizeDocuments(
  documents: LocalDocument[],
  maxIterations = 5,
) {
  const formattedDocs = documents.map(
    (doc) =>
      new Document({
        pageContent: doc.content,
        metadata: {
          title: doc.title,
          link: doc.link,
          date: doc.date,
          source: doc.source,
          selector: doc.selector,
          index: doc.index,
        },
      }),
  );

  let summaries = await runMappers(formattedDocs);

  const shouldCollapse = await checkShouldCollapse(summaries);
  if (shouldCollapse) {
    summaries = await collapseSummaries(summaries, maxIterations);
  }
  const finalSummary = await reduceSummaries(summaries);
  console.log("finalSummary", finalSummary);
}

Conclusions

After implementing our application, we can clearly see how the Map-Reduce pattern applies to the world of generative AI. Using this pattern, we can overcome LLM context window limits. In software development, simple, battle-tested solutions invented many years ago can work perfectly alongside brand-new, groundbreaking technologies like AI. This shows that before adopting any new technology, we should ask what patterns we can use to overcome its limitations. Don’t be afraid to use battle-tested patterns and algorithms, and combine them with code that leverages state-of-the-art technologies.

If you want to learn more about the context window, check out my article on managing the context window of GPT-4o-mini

Also, check out the repository of our blog summarizer on GitHub.

PS: This post was originally published here on my tech blog called about-js

Managing the Context Window of GPT-4o-mini in JavaScript

Grzegorz Dubiel — Tue, 26 Aug 2025 13:37:26 +0000

GPT-4o-mini is a powerful, fast, and affordable language model that can be used for a wide range of tasks in different languages. The context window of this model is 128k tokens.

You might think that this sounds like a lot, but when we compare it with things like an excessive prompt consisting of a complicated schema, a few-shot example, and the content of the document you want to analyze, along with the LLM’s response, we will quickly realize that it is not that much. In this post, we will dive into this topic by building a simple keyword extractor using GPT-4o-mini and Node.js.

Why context window management is important?

Stability and predictability are among the most important aspects here. Even if the model's context is large, the API may have limitations on tokens per minute. Model response accuracy and performance are also crucial. Sending all of your text for analysis at once is not the best idea, as the result will likely be poorer because the model will have to focus on a broader context. Sending more requests in parallel leads to faster analysis and higher-quality responses. With these aspects in mind, we can adjust the size and number of requests sent to the model to meet your application’s requirements and stay within the limits.

What are tokens?

A token is the smallest unit of text data. A word is split into fragments. The process of splitting text is handled by the tokenizer. The way in which the text is tokenized varies depending on the LLM model and its tokenizer. For more explanations about tokens, you can refer to this OpenAI article.

What is the context window?

The context window refers to the number of tokens that the LLM can handle at the same time. It's like memory. Imagine you're writing a summary of a book. You can only process a window of 50 pages in your mind; the previous pages are forgotten. This is how the context window works, especially in the case of the keyword extractor we're going to build. In our case, we will focus on a single text input composed of: prompt, output schema, and text for analysis. This approach requires different strategies than maintaining a conversation(chat), such as prompt-response.

Building a keyword extractor

The goal of our application is to extract keywords from the text, it may be an article, book, web page etc.. We have to be prepared for very long texts, so the context window management is crucial.

Strategies for managing the context window

To effectively fulfill the requirements of our application and ensure its stability, we must pursue the following strategies for managing the context window:

Prompt strategies: We can control the output by including guidelines in the prompt. In our case, we can instruct the LLM to output only keywords with a maximum of 2-3 words and return only the 5 most relevant ones. This is important because the OpenAI API does not allow setting a maximum number of items in the schema for structured output.
Counting tokens: To count tokens, we will have to choose a tokenizer and implement counting logic before sending the text to the LLM. Additionally, we must estimate how many tokens the output will contain.
Chunking the text: Chunking text ensures that we don't exceed the context window. In our case, the need for contextual awareness is less important, so we can use more aggressive chunking strategies—sending more but smaller chunks to the LLM while ensuring that words and sentences are not cut off.
Structured output: Forcing the LLM to output structured data and setting the temperature to 0 for a more concise output also gives us control over the context window, as we can easier estimate how many tokens the output will have.

Let's get to work

The first thing we need to have in place is prompts:

export const systemPrompt = `You are a specialized AI assistant focused on keyword extraction from blog posts.
Your primary role is to identify and extract the most relevant keywords and key phrases that best represent the content.
Always provide keywords in a consistent format and focus on technical accuracy.`;

export const mainPrompt = `Analyze the following blog post and extract the most relevant keywords and key phrases.
Guidelines:
- Extract 5 most important keywords/phrases
- Focus on technical terms, concepts, and specific terminology
- Include both single words and short phrases when relevant (2-3 words)
- Exclude common stop words unless part of a specific term
- Add confident scores to each keyword/phrase
Input text:`;

As I mentioned earlier when talking about prompt strategies, we can explicitly tell the LLM how long the keyword list should be and how long each keyword can be.

We end our main prompt with the text Input text: because we will concatenate the input text with the main prompt later in the code.

Once we have the prompts prepared, it's time to define our schema for the structured output. We will use zod.

import { z } from "zod";

export const keywordsResponseSchema = z.object({
  results: z.array(z.object({ keyword: z.string(), confidence: z.number() })),
});

zod works perfectly with the withStructuredOutput method from the langchain library, which we will use to manage the LLM's API requests. langchain converts the zod schema into the JSON schema accepted by the OpenAI API. Additionally, we can take advantage of the types for our outputs, which are automatically inferred from the schema.

Chunking the text and counting tokens

Now it's time for the first part of the logic: we're going to split the text into smaller parts

In our example, the chunk size is set to 2000 tokens. While the entire context window of GPT-4o-mini is 128k tokens, we could technically set the chunk size close to this value, but a smaller chunk performs better in terms of both performance and output quality.

Here is how we declare our function for this:

async function getChunkedText(content: string, maxChunkSize: number = 2000) {
  /* --- */
}

The first thing we need to do to determine how to split the text is perform calculations. For this, we need to use a tokenizer to encode the text into tokens. In our case, js-tiktoken is the best choice because it is compatible with the model we're currently using.

import { encodingForModel } from "js-tiktoken";

async function getChunkedText(content: string, maxChunkSize: number = 2000) {
  const encoder = encodingForModel("gpt-4o-mini-2024-07-18");

  /* --- */
}

To ensure the best estimations, we need to subtract the tokens used by the prompts, schema, and text from the maxChunkSize parameter. Our schema is in the zodType format, so we need to convert it to a string. For this, we use the zod-to-json-schema package.

import { encodingForModel } from "js-tiktoken";
import { mainPrompt, systemPrompt } from "./prompts";
import zodToJsonSchema from "zod-to-json-schema";
import { keywordsResponseSchema } from "./schema";
import { type ZodType } from "zod";

function getSchemaAsString(schema: ZodType<any>) {
  const jsonSchema = zodToJsonSchema(schema);

  const schemaString = JSON.stringify(jsonSchema);

  return schemaString;
}

async function getChunkedText(content: string, maxChunkSize: number = 2000) {
  const encoder = encodingForModel("gpt-4o-mini-2024-07-18");
  const schemaString = getSchemaAsString(keywordsResponseSchema);

  const calculatedChunkSize =
    maxChunkSize -
    encoder.encode(systemPrompt).length -
    encoder.encode(mainPrompt).length -
    encoder.encode(schemaString).length;

  /* --- */
}

As a result of the calculation, we obtain the final chunk size. We provide two values to the splitter's constructor:

chunkSize: The calculated size of the chunk.
lengthFunction: This function counts the tokens and returns the size of the chunk candidate, enabling the splitter to determine if the chunk fits within the allowed threshold.

For our application, we use RecursiveCharacterTextSplitter, which chunks text from sentences down to words. This means that the algorithm first checks if one or more phrases fit within the chunk size threshold. If they don't, it recursively breaks the text down into smaller parts.

import { encodingForModel, Tiktoken } from "js-tiktoken";
import { mainPrompt, systemPrompt } from "./prompts";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
import zodToJsonSchema from "zod-to-json-schema";
import { keywordsResponseSchema } from "./schema";
import { type ZodType } from "zod";

function getSchemaAsString(schema: ZodType<any>) {
  const jsonSchema = zodToJsonSchema(schema);

  const schemaString = JSON.stringify(jsonSchema);

  return schemaString;
}

async function getChunkedText(content: string, maxChunkSize: number = 2000) {
  const encoder = encodingForModel("gpt-4o-mini-2024-07-18");
  const schemaString = getSchemaAsString(keywordsResponseSchema);

  const calculatedChunkSize =
    maxChunkSize -
    encoder.encode(systemPrompt).length -
    encoder.encode(mainPrompt).length -
    encoder.encode(schemaString).length;

  handleLimitError({ calculatedChunkSize, schemaString }, encoder);

  const textSplitter = new RecursiveCharacterTextSplitter({
    chunkSize: calculatedChunkSize,

    lengthFunction: (text: string) => encoder.encode(text).length,
  });
  const texts = await textSplitter.splitText(content);

  return texts;
}

export default getChunkedText;

Sending a request and handling the response

We have a function for splitting the text in place, so we can start sending requests to the model.

The first thing is the configuration. We use langchain's integration for OpenAI.

It is worth mentioning that by setting the temperature to 0, we can indirectly influence the LLM's response, making the output concise and preventing the model from being too creative.

import { ChatOpenAI } from "@langchain/openai";

const model = new ChatOpenAI({
  model: "gpt-4o-mini-2024-07-18",
  temperature: 0,
  apiKey: process.env.OPENAI_API_KEY,
});

/* --- */

Next, we declare a function for calling the LLM. We need to preprocess our data and prompts, so we load a document with the text and chunk it using the function we prepared earlier.

For the purpose of our example, we simply load a file with the text, but the text could also be passed to the function via arguments.

export async function callLlm() {
  const text = loadTextFromFile();
  const textChunks = await getChunkedText(text);
  /*  ---  */
}

Now we can register our zod schema to ensure structured output from the LLM and send the requests. We can take advantage of chunking by running them in parallel and awaiting the responses using Promise.all. This allows us to process the data faster and in smaller parts, helping the LLM focus more on details, which leads to better results. However, remember that when sending requests this way, you need to account for API rate limits.

export async function callLlm() {
  const text = loadTextFromFile();
  const textChunks = await getChunkedText(text);
  const structured = model.withStructuredOutput(keywordsResponseSchema);
  const completions = await Promise.all(
    textChunks.map(async (chunk) => {
      return structured.invoke([
        {
          role: "system",
          content: systemPrompt,
        },
        {
          role: "user",
          content: mainPrompt + chunk,
        },
      ]);
    }),
  );

  /* --- */
}

The last step is output formatting and data filtering. We merge responses from each chunk and filter them based on the confidence level to ensure that only relevant keywords are returned. Finally, we remove duplicates using new Set().

export async function callLlm() {
  const text = loadTextFromFile();
  const textChunks = await getChunkedText(text);
  const structured = model.withStructuredOutput(keywordsResponseSchema);
  const completions = await Promise.all(
    textChunks.map(async (chunk, index) => {
      return structured.invoke([
        {
          role: "system",
          content: systemPrompt,
        },
        {
          role: "user",
          content: mainPrompt + chunk,
        },
      ]);
    }),
  );

  const results = completions
    .flatMap((completion) => completion.results)
    .filter((result) => result.confidence >= 0.85)
    .map((result) => result.keyword);
  return [...new Set(results)];
}

Final response:

Wrapping up

Managing the context window is crucial, especially when focusing on large text analysis. Rate limits can be tricky, and despite the model having a 128k token window, these limits can prevent full utilization. For example, currently, for GPT-4o-mini hosted on Azure Open AI, the max initial quota is 30k tokens per minute! To get more, you have to request a larger quota from Microsoft. Even if your limits are higher, maintaining control over the context window is essential, whether for a production-grade application or a fun project. No one wants to run their AI integrations while constantly encountering errors.

Thanks for reading 🙌

PS: Look at the repo of the keywords extractor app here on GitHub

PS2: This post was originally published here on my tech blog called about-js

Async Local Storage is Here to Help You

Grzegorz Dubiel — Wed, 18 Dec 2024 00:38:21 +0000

When you hear the phrase "Async Local Storage," what comes to mind? You might initially think it refers to some magical implementation of browser-based local storage. However, this assumption is incorrect. Async Local Storage is neither browser-related nor a typical storage mechanism. Probably one or two libraries you have used use it under the hood. In many cases, this feature can save you from dealing with messy code.

What is Async Local Storage?

Async Local Storage is a feature introduced in Node.js, initially added in versions v13.10.0 and v12.17.0, and later stabilized in v16.4.0. It is part of the async_hooks module, which provides a way to track asynchronous resources in Node.js applications. The feature enables the creation of a shared context that multiple asynchronous functions can access without explicitly passing it. The context is available in every (and only) operation executed within the callback passed to the run() method of the AsyncLocalStorage instance.

A pattern for using AsyncLocalStorage

Before diving into the examples, let’s explain the pattern we will be using.

Initialization

import { AsyncLocalStorage } from "async_hooks";
import { Context } from "./types";

export const asyncLocalStorage = new AsyncLocalStorage<Context>();

// export const authAsyncLocalStorage = new AuthAsyncLocalStorage<AuthContext>()

In the module above, we initialize an instance of AsyncLocalStorage and export it as a variable.

Usage

asyncLocalStorage.run({ userId }, async () => {
  const usersData: UserData = await collectUsersData();
  console.log("usersData", usersData);
});

// (method) AsyncLocalStorage<unknown>.run<Promise<void>>(store: unknown, callback: () => Promise<void>): Promise<void> (+1 overload)

The run() method takes two arguments: storage, which contains the data we want to share, and callback, where we place our logic. As a result, the storage becomes accessible in every function call within the callback, allowing for seamless data sharing across asynchronous operations.

async function collectUsersData() {
  const context = asyncLocalStorage.getStore();
}

To access the context, we import our instance and call the asyncLocalStorage.getStore() method. The great thing is that the storage retrieved from getStore() is typed because we passed the Context type to AsyncLocalStorage during initialization: new AsyncLocalStorage<Context>().

Async Local Storage as an auth context

There is no web application without an authentication system. We must validate auth tokens and extract user information. Once we obtain the user identity, we want to make it available in the route handlers and avoid duplicating code in each one. Let’s see how we can utilize AsyncLocalStorage to implement an auth context while keeping our code clean.

I chose fastify for this example.

According to the documentation fastify is:

Fast and low overhead web framework, for Node.js

Ok, let's get started:

Install fastify:

npm install fastify

Define type for our auth context:

type Context = Map<"userId", string>;

Initialize an instance of AsyncLocalStorage, assign it to a variable, and export the variable. Remember to pass the relevant type: new AsyncLocalStorage<Context>().

import { AsyncLocalStorage } from "async_hooks";
import { Context } from "./types";

export const authAsyncLocalStorage = new AsyncLocalStorage<Context>();

Initialize a Fastify instance and add a utility for error handling:

import Fastify from "fastify";

/* other code... */

const app = Fastify();

function sendUnauthorized(reply: FastifyReply, message: string) {
  reply.code(401).send({ error: `Unauthorized: ${message}` });
}

/* other code... */

Now comes the very important part. We are going to add an onRequest hook to wrap handlers with the authAsyncLocalStorage.run() method.

import Fastify from "fastify";
import { authAsyncLocalStorage } from "./context";
import { getUserIdFromToken, validateToken } from "./utils";

/* other code... */

app.addHook(
  "onRequest",
  (request: FastifyRequest, reply: FastifyReply, done: () => void) => {
    const accessToken = request.headers.authorization?.split(" ")[1];
    const isTokenValid = validateToken(accessToken);
    if (!isTokenValid) {
      sendUnauthorized(reply, "Access token is invalid");
    }
    const userId = accessToken ? getUserIdFromToken(accessToken) : null;

    if (!userId) {
      sendUnauthorized(reply, "Invalid or expired token");
    }
    authAsyncLocalStorage.run(new Map([["userId", userId]]), async () => {
      await new Promise((resolve) => setTimeout(resolve, 2000));
      sendUnauthorized(reply, "Invalid or expired token");
      done();
    });
  },
);

/* other code... */

After successful validation, we call the run() method from our authAsyncLocalStorage. As the storage argument, we pass the auth context with the userId retrieved from the token. In the callback, we call the done function to continue with the Fastify lifecycle.

If we have authentication checks that require asynchronous operations, we should add them to the callback. This is because, according to the documentation:

the done callback is not available when using async/await or returning a Promise. If you do invoke a done callback in this situation unexpected behavior may occur, e.g. duplicate invocation of handlers

Here's an example of how that might look:

/* other code... */

authAsyncLocalStorage.run(new Map([["userId", userId]]), async () => {
  const isUserValid = await checkIfUserIsValid(userId);

  if (!isUserValid) {
    sendUnauthorized(reply, "User identity not valid");
  }
  done();
});

Our example has only one protected route. In more complex scenarios, you might need to wrap only specific routes with the authentication context. In such cases, you could either:

Wrap the onRequest hook in a custom plugin that's applied only to specific routes.
Add route distinction logic within the onRequest hook itself.

All right, our context is set and we can now define a protected route:

import { UserRepository } from "./user-repository";
import { getContext } from "./with-async-local-storage";
import { Context } from "./types";

/* other code... */

app.get("/email-addresses", async () => {
  const context = authAsyncLocalStorage.getStore();
  const userId = context.get("userId");
  const userRepository = new UserRepository();
  const addresses = await userRepository.getEmailAddresses(userId);

  return { addresses };
});

The code is pretty straightforward. We import authAsyncLocalStorage, retrieve the userId, initialize UserRepository, and fetch data. This approach keeps the route handler clean and focused.

Exploring How Next.js uses Async Local Storage

In this example, we'll reimplement the cookies helper from Next.js. But wait—this is a post about AsyncLocalStorage, right? So why are we talking about cookies? The answer is simple: Next.js uses AsyncLocalStorage to manage cookies on the server. That's why reading a cookie in a server component is as easy as:

import { cookies } from "next/headers";

export default function ExamplePage() {
  const cookieStore = await cookies();
  const testCookie = cookieStore.get("test");

  return (
    <div>
      <h1>Server Component: Cookie Usage</h1>
      {testCookie?.value || ""}
    </div>
  );
}

We use the cookies function exported from next/headers, which provides several methods for managing cookies. But how is this technically possible?

Now it's time to start our re-implementation"

First, I want to mention that this example is based on the knowledge I gained from a great video, by Lee Robinson and diving in the Next.js repository.

In this example, we'll use Hono as our server framework. I chose it for two reasons:

I just wanted to give it a try.
It offers solid support for JSX.

First install Hono:

npm install hono

Now, initialize Hono and add middleware:

import { serve } from "@hono/node-server";
import { Hono } from "hono";
import { logger } from "hono/logger";
import withAsyncLocalStorage from "./with-async-local-storage";
import { setCookieContext } from "./cookies";

const app = new Hono();

app.use(logger());

app.use(async (c, next) => {
  return cookieAsyncLocalStorage.run(setCookieContext(c), async () => {
    await next();
  });
});

serve(
  {
    fetch: app.fetch,
    port: 3000,
  },
  (info) => {
    console.log(`Server is running on http://localhost:${info.port}`);
  },
);

The code resembles the middleware from the Fastify example, does it? To set the context, we utilize setCookieContext, which is imported from the cookies module — our custom simple implementation of the cookies function. Let's follow the setCookieContext function and navigate to the module from which it was imported:

import type { Context, Env } from "hono";
import { deleteCookie, getCookie, setCookie } from "hono/cookie";
import { cookieAsyncLocalStorage } from "./context";

const MOCK_ERROR = "Can not set or delete cookie in server component";

export function setCookieContext(c: Context<Env, never, {}>) {
  return {
    cookies: getCookie(c),
    setCookie: (key: string, value: string) => setCookie(c, key, value),
    deleteCookie: (key: string) => deleteCookie(c, key),
  };
}

function cookies() {
  const context = cookieAsyncLocalStorage.getStore();

  const methods = {
    getCookies: () => {
      if (!context?.cookies) return [];
      return Object.entries(context.cookies).map(([key, value]) => ({
        key,
        value,
      }));
    },
    setCookie: (key: string, value: string) => {
      throw new Error(MOCK_ERROR);
    },
    deleteCookie: (key: string) => {
      throw new Error(MOCK_ERROR);
    },
  };

  return new Promise<typeof methods>((resolve) => {
    resolve(methods);
  });
}

export default cookies;

The setCookieContext function (whose return value we passed to cookieAsyncLocalStorage.run() in the Hono middleware) extracts cookies from the c parameter, which represents the hono context, and bundles them with closures that provide utility functions for managing cookies.

Our cookies function replicates the functionality of the cookies from next/headers. It utilizes the cookieAsyncLocalStorage.getStore() method to access the same context that is passed to cookieAsyncLocalStorage.run() when it is called.

We wrapped the return of our cookies function in a promise to mimic the behavior of the Next.js implementation. Prior to version 15, this function was synchronous. Now, in the current Next.js code, the methods returned by the cookies are attached to a promise object, as shown in the following simplified example:

/*  .......... */
// https://github.com/vercel/next.js/blob/canary/packages/next/src/server/request/cookies.ts

Object.defineProperties(promise, {
  [Symbol.iterator]: {
    get: {
      value: function () {
        /* .... */
      },
    },
    getAll: {
      value: function () {
        /* .... */
      },
    },
    has: {
      value: function () {
        /* .... */
      },
    },
    set: {
      value: function () {
        /* .... */
      },
    },
    delete: {
      value: function () {
        /* .... */
      },
    },
    /* .... */
  },
});

/* .... */

Another point worth mentioning is that in our case, using cookies.setCookie and cookies.deleteCookie always throws an error, similar to the behavior observed in Next.js when setting cookies in a server component. We hardcoded this logic because, in the original implementation, whether we can use setCookie or deleteCookie depends on the phase(WorkUnitPhase) property stored in the storage called RequestStore(this is the implementation of AsyncLocalStorage and also stores cookies). However, that topic is better suited for another post. To keep this example simple, let's omit the simulation of WorkUnitPhase.

Now we need to add our React code.

Add the App component:

import type { FC } from "hono/jsx";

const App: FC = ({ children }) => {
    return (
        <html>
            <body>{children}</body>
        </html>
    );
};

export default App;

Add a component for managing cookies:

import { FC } from "hono/jsx";
import cookies from "../cookies";

const DisplayCookies: FC = async () => {
    const cookieStore = await cookies();

    return (
        <div>
            <h1>Hello!</h1>
            <p>Here are your cookies:</p>
            <ul>
                {cookieStore.getCookies().map((cookie) => (
                    <li key={cookie.key}>
                        {cookie.key}: {cookie.value}
                    </li>
                ))}
            </ul>
        </div>
    );
};

export default DisplayCookies;

The usage of cookies is similar to how it is used in Next.js React server components.

Add a route handler to render the template:

import { serve } from "@hono/node-server";
import { Hono } from "hono";
import { logger } from "hono/logger";
import App from "./components/app";
import DisplayCookies from "./components/display-cookies";
import { setCookieContext } from "./cookies";
import { cookieAsyncLocalStorage } from "./context";

const app = new Hono();

app.use(logger());

app.use(async (c, next) => {
  return cookieAsyncLocalStorage.run(setCookieContext(c), async () => {
    await next();
  });
});

app.get("/", async (c) => {
  const renderDisplayCookiesComponent = await (<DisplayCookies />);
  return c.html(<App>{renderDisplayCookiesComponent}</App>);
});

serve(
  {
    fetch: app.fetch,
    port: 8000,
  },
  (info) => {
    console.log(`Server is running on http://localhost:${info.port}`);
  },
);

Our template is rendered by the html method from the hono context. The key point here is that the route handler runs within the asyncLocalStorage.run() method, which takes cookieContext. As a result, we can access this context in the DisplayCookies component through the cookies function.

It is not possible to set cookies inside React server components, so we need to do it manually:

Let's refresh a page:

And here we are, our cookies are successfully retrieved and displayed.

Conclusions

There are many more use cases for asyncLocalStorage. This feature allows you to build custom contexts in nearly any server framework. The asyncLocalStorage context is encapsulated within the execution of the run() method, making it easy to manage. It’s perfect for handling request-based scenarios. The API is simple and flexible, enabling scalability by creating instances for each state. It is possible to seamlessly maintain separate contexts for things like authentication, logging, and feature flags.

Despite its benefits, there are a few considerations to keep in mind. I've heard opinions that asyncLocalStorage introduces too much 'magic' into the code. I'll admit that when I first used this feature, it took me some time to fully grasp the concept. Another thing to consider is that importing the context into a module creates a new dependency that you’ll need to manage. However, in the end, passing values through deeply nested function calls is much worse.

Thanks for reading, and see you in the next post!👋

PS: You can find the examples(plus one bonus) here

Bog Post source: https://www.aboutjs.dev/en/async-local-storage-is-here-to-help-you