Fix llms.txt not including all content (#4260)

Author: gregberge

.changeset/fix-llms-full-sections-pagination.md

@@ -0,0 +1,5 @@
+---
+"gitbook": patch
+---
+
+Fix llms-full.txt pagination to include pages from all site sections.

packages/gitbook/src/routes/llms-full.ts

@@ -10,7 +10,7 @@ import { filterSiteSpacesByLocale, getSiteStructureSections } from '@/lib/sites'
 import type { RevisionPageDocument, SiteSection, SiteSpace } from '@gitbook/api';
 import assertNever from 'assert-never';
 import type { Paragraph } from 'mdast';
-import { pMapIterable } from 'p-map';
+import pMap, { pMapIterable } from 'p-map';
 
 // We limit the concurrency to 100 to avoid reaching limit with concurrent requests
 // or file descriptor limits.
@@ -19,6 +19,8 @@ const MAX_CONCURRENCY = 100;
 // Default limit for pages per batch
 const DEFAULT_PAGE_LIMIT = 100;
 
+type MarkdownPageEntry = { context: GitBookSiteContext; page: RevisionPageDocument };
+
 /**
  * Generate a llms-full.txt file for the site.
  * As the result can be large, we stream it as we generate it.
@@ -29,11 +31,16 @@ export async function serveLLMsFullTxt(context: GitBookSiteContext, page = 0) {
     }
 
     const offset = page * DEFAULT_PAGE_LIMIT;
+    const allPages = await getMarkdownPageEntriesFromSiteStructure(context);
+
+    if (allPages.length <= offset) {
+        return new Response('No content found', { status: 404 });
+    }
 
     return new Response(
         new ReadableStream<Uint8Array>({
             async pull(controller) {
-                await streamMarkdownFromSiteStructure(context, controller, offset);
+                await streamMarkdownPageEntries(context, controller, allPages, offset);
                 controller.close();
             },
         }),
@@ -46,100 +53,100 @@ export async function serveLLMsFullTxt(context: GitBookSiteContext, page = 0) {
 }
 
 /**
- * Stream markdown from site structure.
+ * Get the document pages that should be included in the full llms.txt output for a site structure.
  */
-async function streamMarkdownFromSiteStructure(
-    context: GitBookSiteContext,
-    stream: ReadableStreamDefaultController<Uint8Array>,
-    offset: number
-): Promise<void> {
+async function getMarkdownPageEntriesFromSiteStructure(
+    context: GitBookSiteContext
+): Promise<MarkdownPageEntry[]> {
     switch (context.structure.type) {
         case 'sections':
-            return streamMarkdownFromSections(
+            return getMarkdownPageEntriesFromSections(
                 context,
-                stream,
-                getSiteStructureSections(context.structure, { ignoreGroups: true }),
-                offset
+                getSiteStructureSections(context.structure, { ignoreGroups: true })
             );
         case 'siteSpaces':
-            await streamMarkdownFromSiteSpaces(
-                context,
-                stream,
-                context.structure.structure,
-                offset
-            );
-            return;
+            return getMarkdownPageEntriesFromSiteSpaces(context, context.structure.structure);
         default:
             assertNever(context.structure);
     }
 }
 
 /**
- * Stream markdown from site sections.
+ * Get the document pages that should be included in the full llms.txt output for site sections.
  */
-async function streamMarkdownFromSections(
+async function getMarkdownPageEntriesFromSections(
     context: GitBookSiteContext,
-    stream: ReadableStreamDefaultController<Uint8Array>,
-    siteSections: SiteSection[],
-    offset: number
-): Promise<void> {
-    let currentPageIndex = 0;
-
-    for (const siteSection of siteSections) {
-        const result = await streamMarkdownFromSiteSpaces(
-            context,
-            stream,
-            siteSection.siteSpaces,
-            offset,
-            currentPageIndex
-        );
-        currentPageIndex = result.currentPageIndex;
-
-        if (result.reachedLimit) {
-            break;
-        }
-    }
+    siteSections: SiteSection[]
+): Promise<MarkdownPageEntry[]> {
+    return getMarkdownPageEntriesFromFilteredSiteSpaces(
+        siteSections.flatMap((siteSection) =>
+            filterSiteSpacesByLocale(siteSection.siteSpaces, context.locale)
+        ),
+        context
+    );
 }
 
 /**
- * Stream markdown from site spaces.
+ * Get the document pages that should be included in the full llms.txt output for site spaces.
  */
-export async function streamMarkdownFromSiteSpaces(
+async function getMarkdownPageEntriesFromSiteSpaces(
     context: GitBookSiteContext,
-    stream: ReadableStreamDefaultController<Uint8Array>,
-    siteSpaces: SiteSpace[],
-    offset = 0,
-    initialPageIndex = 0
-): Promise<{ currentPageIndex: number; reachedLimit: boolean }> {
-    let totalPagesProcessed = initialPageIndex;
-
-    // Collect all pages first
-    const allPages: Array<{ context: GitBookSiteContext; page: RevisionPageDocument }> = [];
-
-    const filteredSiteSpaces = filterSiteSpacesByLocale(siteSpaces, context.locale);
+    siteSpaces: SiteSpace[]
+): Promise<MarkdownPageEntry[]> {
+    return getMarkdownPageEntriesFromFilteredSiteSpaces(
+        filterSiteSpacesByLocale(siteSpaces, context.locale),
+        context
+    );
+}
 
-    for (const siteSpace of filteredSiteSpaces) {
-        const siteSpaceUrl = siteSpace.urls.published;
-        if (!siteSpaceUrl) {
-            continue;
-        }
-        const siteSpaceContext = await fetchSiteContextForSiteSpace(context, siteSpace);
-        const pages = getIndexablePages(siteSpaceContext.revision.pages);
-
-        // Add document pages to our collection
-        for (const { page } of pages) {
-            if (page.type === 'document') {
-                allPages.push({
-                    context: siteSpaceContext,
-                    page,
-                });
-            }
+/**
+ * Get markdown page entries from already locale-filtered site spaces.
+ */
+async function getMarkdownPageEntriesFromFilteredSiteSpaces(
+    siteSpaces: SiteSpace[],
+    context: GitBookSiteContext
+): Promise<MarkdownPageEntry[]> {
+    const publishedSiteSpaces = siteSpaces.filter((siteSpace) => siteSpace.urls.published);
+
+    const allPages = await pMap(
+        publishedSiteSpaces,
+        async (siteSpace): Promise<MarkdownPageEntry[]> => {
+            const siteSpaceContext = await fetchSiteContextForSiteSpace(context, siteSpace);
+            const pages = getIndexablePages(siteSpaceContext.revision.pages);
+
+            return pages.flatMap(({ page }) => {
+                if (page.type !== 'document') {
+                    return [];
+                }
+
+                return [
+                    {
+                        context: siteSpaceContext,
+                        page,
+                    },
+                ];
+            });
+        },
+        {
+            concurrency: MAX_CONCURRENCY,
         }
-    }
+    );
+
+    return allPages.flat();
+}
 
+/**
+ * Stream a single paginated window of markdown page entries.
+ */
+async function streamMarkdownPageEntries(
+    context: GitBookSiteContext,
+    stream: ReadableStreamDefaultController<Uint8Array>,
+    allPages: MarkdownPageEntry[],
+    offset: number
+): Promise<{ currentPageIndex: number; reachedLimit: boolean }> {
     // Apply pagination - skip pages before offset
     const pagesToProcess = allPages.slice(offset, offset + DEFAULT_PAGE_LIMIT);
-    totalPagesProcessed = offset;
+    let totalPagesProcessed = offset;
 
     // Process the pages
     for await (const markdown of pMapIterable(

packages/gitbook/tests/llms.test.ts

@@ -116,4 +116,16 @@ describe('llms-full.txt', () => {
         },
         { timeout: 30_000 }
     );
+
+    it(
+        'should return 404 when a llms-full.txt page has no content',
+        async () => {
+            const response = await fetch(
+                getContentTestURL('https://gitbook.gitbook.io/test-gitbook-open/llms-full.txt/999')
+            );
+
+            expect(response.status).toBe(404);
+        },
+        { timeout: 30_000 }
+    );
 });