fix(knowledge): bound JSON/YAML chunker expansion (#7524)

* fix(knowledge): bound JSON/YAML chunker expansion

JsonYamlChunker re-parsed and re-serialized document content with no
expansion limit. ChunkBudget only counts emitted chunks, so every parse
and full-object stringify ran before it could fire — a small aliased YAML
source expands to tens of MB, and the same content was parsed twice
because isStructuredData and chunkJsonYaml each parsed it.

- Measure the parsed value with the shared measureYamlExpansion guard
  before anything materializes it, and skip parsing entirely when the
  source is already larger than the ceiling
- Size the ceiling to the most text the chunker could ever emit
  (maxChunks x chunkSize), floored at 4MB and capped at what the YAML
  file parser itself permits, so documents that fit the budget chunk
  exactly as before
- Replace isStructuredData + chunkJsonYaml with one chunkStructured entry
  point that parses once and returns null when the content is not
  structured, leaving chunker selection with the document processor

* fix(knowledge): allow proportionate expansion in the chunker guard

Comparing the expansion estimate against the output budget mixed two
units: measureYamlExpansion charges a flat per-node allowance, so a
document of many small values is charged several times its pretty-printed
size and was rejected even though its chunks fit the budget — a flat array
of a million booleans is charged ~22MB against ~8MB of real output.

Allow the larger of two admissible expansions: one that fits the output
budget, and one proportionate to the source. Alias expansion overshoots
its source by orders of magnitude, so it is still rejected, while an
ordinary large document is no longer charged for the estimator's
conservatism. Neither allowance ever exceeds the file parser's own cap.

Author: waleedlatif1

apps/sim/lib/chunkers/json-yaml-chunker.test.ts

@@ -3,7 +3,7 @@
  */
 
 import { describe, expect, it, vi } from 'vitest'
-import { JsonYamlChunker } from './json-yaml-chunker'
+import { JsonYamlChunker } from '@/lib/chunkers/json-yaml-chunker'
 
 vi.mock('@/lib/tokenization', () => ({
   getAccurateTokenCount: (text: string) => Math.ceil(text.length / 4),
@@ -37,30 +37,110 @@ describe('JsonYamlChunker', () => {
     expect(chunks.every((chunk) => chunk.tokenCount <= 100)).toBe(true)
   })
 
-  describe('isStructuredData', () => {
-    it('should detect valid JSON', () => {
-      expect(JsonYamlChunker.isStructuredData('{"key": "value"}')).toBe(true)
+  describe('chunkStructured', () => {
+    it('chunks valid JSON', async () => {
+      await expect(JsonYamlChunker.chunkStructured('{"key": "value"}')).resolves.not.toBeNull()
     })
 
-    it('should detect valid JSON array', () => {
-      expect(JsonYamlChunker.isStructuredData('[1, 2, 3]')).toBe(true)
+    it('chunks a valid JSON array', async () => {
+      await expect(JsonYamlChunker.chunkStructured('[1, 2, 3]')).resolves.not.toBeNull()
     })
 
-    it('should detect valid YAML', () => {
-      expect(JsonYamlChunker.isStructuredData('key: value\nother: data')).toBe(true)
+    it('chunks valid YAML', async () => {
+      await expect(
+        JsonYamlChunker.chunkStructured('key: value\nother: data')
+      ).resolves.not.toBeNull()
     })
 
-    it('should return false for plain text parsed as YAML scalar', () => {
-      expect(JsonYamlChunker.isStructuredData('Hello, this is plain text.')).toBe(false)
+    it('declines plain text that parses as a YAML scalar', async () => {
+      await expect(
+        JsonYamlChunker.chunkStructured('Hello, this is plain text.')
+      ).resolves.toBeNull()
     })
 
-    it('should return false for invalid JSON/YAML with unbalanced braces', () => {
-      expect(JsonYamlChunker.isStructuredData('{invalid: json: content: {{')).toBe(false)
+    it('declines invalid JSON/YAML with unbalanced braces', async () => {
+      await expect(
+        JsonYamlChunker.chunkStructured('{invalid: json: content: {{')
+      ).resolves.toBeNull()
     })
 
-    it('should detect nested JSON objects', () => {
+    it('chunks nested JSON objects', async () => {
       const nested = JSON.stringify({ level1: { level2: { level3: 'value' } } })
-      expect(JsonYamlChunker.isStructuredData(nested)).toBe(true)
+      await expect(JsonYamlChunker.chunkStructured(nested)).resolves.not.toBeNull()
+    })
+
+    it('declines an alias-expansion bomb instead of expanding it', async () => {
+      const lines = ['a0: &a0 "lol"']
+      for (let level = 1; level <= 7; level++) {
+        lines.push(
+          `a${level}: &a${level} [${Array(7)
+            .fill(`*a${level - 1}`)
+            .join(',')}]`
+        )
+      }
+      lines.push('top: *a7')
+      const bomb = lines.join('\n')
+
+      const chunks = await JsonYamlChunker.chunkStructured(bomb, {
+        chunkSize: 1024,
+        minCharactersPerChunk: 1,
+        maxChunks: 5000,
+      })
+
+      expect(chunks).toBeNull()
+    })
+
+    it('keeps structure for a many-small-node document that fits the budget', async () => {
+      const flags = JSON.stringify(Array.from({ length: 250_000 }, (_, i) => i % 2 === 0))
+
+      const chunks = await JsonYamlChunker.chunkStructured(flags, {
+        chunkSize: 1024,
+        minCharactersPerChunk: 1,
+        maxChunks: 1000,
+      })
+
+      expect(chunks).not.toBeNull()
+      expect(chunks?.length).toBeGreaterThan(1)
+      expect(chunks?.[0].text).toContain('true')
+    })
+
+    it('never parses source larger than one output budget', async () => {
+      const oversized = JSON.stringify({ value: 'x'.repeat(5 * 1024 * 1024) })
+      const parse = vi.spyOn(JSON, 'parse')
+
+      try {
+        await expect(
+          JsonYamlChunker.chunkStructured(oversized, {
+            chunkSize: 1024,
+            minCharactersPerChunk: 1,
+            maxChunks: 1024,
+          })
+        ).resolves.toBeNull()
+        expect(parse).not.toHaveBeenCalled()
+      } finally {
+        parse.mockRestore()
+      }
+    })
+
+    it('chunks with default options', async () => {
+      const chunks = await JsonYamlChunker.chunkStructured(JSON.stringify({ test: 'value' }))
+
+      expect(chunks?.length).toBeGreaterThan(0)
+    })
+
+    it('honors a custom chunk size', async () => {
+      const largeObject: Record<string, string> = {}
+      for (let i = 0; i < 50; i++) {
+        largeObject[`key${i}`] = `value${i}`.repeat(20)
+      }
+      const json = JSON.stringify(largeObject)
+
+      const chunksSmall = await JsonYamlChunker.chunkStructured(json, { chunkSize: 50 })
+      const chunksLarge = await JsonYamlChunker.chunkStructured(json, { chunkSize: 500 })
+
+      expect(chunksSmall).not.toBeNull()
+      expect(chunksLarge).not.toBeNull()
+      expect(chunksSmall?.length).toBeGreaterThan(chunksLarge?.length as number)
     })
   })
 
@@ -368,28 +448,6 @@ server:
     })
   })
 
-  describe('static chunkJsonYaml method', () => {
-    it.concurrent('should work with default options', async () => {
-      const json = JSON.stringify({ test: 'value' })
-      const chunks = await JsonYamlChunker.chunkJsonYaml(json)
-
-      expect(chunks.length).toBeGreaterThan(0)
-    })
-
-    it.concurrent('should accept custom options', async () => {
-      const largeObject: Record<string, string> = {}
-      for (let i = 0; i < 50; i++) {
-        largeObject[`key${i}`] = `value${i}`.repeat(20)
-      }
-      const json = JSON.stringify(largeObject)
-
-      const chunksSmall = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 50 })
-      const chunksLarge = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 500 })
-
-      expect(chunksSmall.length).toBeGreaterThan(chunksLarge.length)
-    })
-  })
-
   describe('chunk metadata', () => {
     it('preserves every source character and offset when bounding oversized chunks', async () => {
       const key = 'p'.repeat(80)

apps/sim/lib/chunkers/json-yaml-chunker.ts

@@ -9,6 +9,8 @@ import {
   normalizeTokenChunkSize,
   tokensToChars,
 } from '@/lib/chunkers/utils'
+import { measureYamlExpansion, type YamlExpansionLimits } from '@/lib/file-parsers/yaml-limits'
+import { FILE_PARSER_YAML_LIMITS } from '@/lib/file-parsers/yaml-parser'
 
 const logger = createLogger('JsonYamlChunker')
 
@@ -20,40 +22,148 @@ type BoundedChunkMetadataMode = 'text-offsets' | 'preserve-range'
 
 const MAX_DEPTH = 5
 
+/**
+ * Smallest source ceiling this chunker imposes, so a knowledge base configured
+ * with tiny chunks keeps structural chunking on documents it indexes perfectly
+ * well today.
+ */
+const MIN_SOURCE_BYTES = 4 * 1024 * 1024
+
+/**
+ * How far a document may legitimately expand past its own source.
+ *
+ * `measureYamlExpansion` charges a flat per-node allowance, so a compact source
+ * of small values is charged well above its own length — `[1,1,1]` costs about
+ * 22 estimated bytes per element against two in source. An order of magnitude of
+ * headroom therefore covers ordinary document shape, while alias expansion
+ * overshoots it by several orders.
+ */
+const MAX_EXPANSION_RATIO = 16
+
+/**
+ * Longest source this chunker will parse: the most text it could ever emit, one
+ * output budget's worth. A larger document cannot be indexed whole by any
+ * chunker — `ChunkBudget` stops it either way — so parsing it buys nothing.
+ */
+function resolveMaxSourceBytes(maxChunks: number | undefined, chunkSize: number): number {
+  if (maxChunks === undefined) return FILE_PARSER_YAML_LIMITS.maxSerializedBytes
+
+  return Math.min(
+    FILE_PARSER_YAML_LIMITS.maxSerializedBytes,
+    Math.max(MIN_SOURCE_BYTES, maxChunks * tokensToChars(chunkSize))
+  )
+}
+
+/**
+ * What the document is allowed to expand to once it is walked as a tree.
+ *
+ * Structural chunking re-serializes what it parsed, so its cost follows the
+ * document's *expanded* size rather than its source size, and `yaml.load`
+ * resolves aliases into shared references — a sub-kilobyte source can carry tens
+ * of megabytes of expansion. `ChunkBudget` cannot bound that: it counts emitted
+ * chunks, and every parse and serialization happens before the first is emitted.
+ *
+ * Two expansions are admissible: one that stays within the output budget, and
+ * one that stays proportionate to the source. Taking the larger of the two keeps
+ * transient allocation tied to work the chunker would have done anyway, without
+ * charging an ordinary large document for the estimator's per-node conservatism.
+ * Neither is ever allowed past what the file parser itself would hand over.
+ */
+function resolveExpansionLimits(sourceBytes: number, maxSourceBytes: number): YamlExpansionLimits {
+  return {
+    /** Bytes bind here; every reached node charges some, so a self-referential anchor still terminates. */
+    maxNodes: Number.MAX_SAFE_INTEGER,
+    maxSerializedBytes: Math.min(
+      FILE_PARSER_YAML_LIMITS.maxSerializedBytes,
+      Math.max(maxSourceBytes, sourceBytes * MAX_EXPANSION_RATIO)
+    ),
+    maxDepth: FILE_PARSER_YAML_LIMITS.maxDepth,
+  }
+}
+
 export class JsonYamlChunker {
   private chunkSize: number
   private minCharactersPerChunk: number
   private maxChunks?: number
+  private readonly maxSourceBytes: number
 
   constructor(options: ChunkerOptions = {}) {
     this.chunkSize = normalizeTokenChunkSize(options.chunkSize ?? 1024, 'JSON/YAML chunk size')
     this.minCharactersPerChunk = options.minCharactersPerChunk ?? 100
     this.maxChunks = options.maxChunks
+    this.maxSourceBytes = resolveMaxSourceBytes(this.maxChunks, this.chunkSize)
   }
 
-  static isStructuredData(content: string): boolean {
+  /**
+   * Read `content` as JSON, falling back to YAML, and measure what the parsed
+   * value expands to before anything materializes it.
+   *
+   * The source-length check comes first so oversized content is never parsed at
+   * all; the expansion measurement then catches what length alone cannot — alias
+   * expansion, and the indentation a pretty-printed re-serialization adds.
+   */
+  private parseWithinLimits(content: string): JsonValue | undefined {
+    if (content.length > this.maxSourceBytes) {
+      return this.reject(
+        `source of ${content.length} characters exceeds the ${this.maxSourceBytes}-byte ceiling`
+      )
+    }
+
+    let parsed: unknown
     try {
-      const parsed = JSON.parse(content)
-      return typeof parsed === 'object' && parsed !== null
+      parsed = JSON.parse(content)
     } catch {
       try {
-        const parsed = yaml.load(content)
-        return typeof parsed === 'object' && parsed !== null
+        parsed = yaml.load(content)
       } catch {
-        return false
+        return undefined
       }
     }
+
+    if (parsed === undefined) return undefined
+
+    const limits = resolveExpansionLimits(content.length, this.maxSourceBytes)
+    const measured = measureYamlExpansion(parsed, limits)
+    if (!measured.within) return this.reject(measured.reason)
+
+    return parsed as JsonValue
   }
 
-  async chunk(content: string): Promise<Chunk[]> {
-    try {
-      let data: JsonValue
-      try {
-        data = JSON.parse(content) as JsonValue
-      } catch {
-        data = yaml.load(content) as JsonValue
+  private reject(reason: string): undefined {
+    logger.warn(
+      'Structured content exceeds the chunking expansion limits, declining to expand it',
+      {
+        reason,
       }
+    )
+    return undefined
+  }
+
+  /**
+   * Chunk `content` as a structured object or array, or return `null` when it is
+   * neither — including when its expanded form outgrows the ceiling above. The
+   * caller then chooses another chunker for it.
+   */
+  static async chunkStructured(
+    content: string,
+    options: ChunkerOptions = {}
+  ): Promise<Chunk[] | null> {
+    const chunker = new JsonYamlChunker(options)
+    const data = chunker.parseWithinLimits(content)
+    if (data === null || typeof data !== 'object') return null
+
+    return chunker.chunkParsed(data, content)
+  }
+
+  async chunk(content: string): Promise<Chunk[]> {
+    const data = this.parseWithinLimits(content)
+    if (data === undefined) return this.chunkAsText(content)
+
+    return this.chunkParsed(data, content)
+  }
 
+  private chunkParsed(data: JsonValue, content: string): Chunk[] {
+    try {
       const chunks: Chunk[] = []
       this.chunkStructuredData(data, [], 0, chunks, new ChunkBudget(this.maxChunks))
 
@@ -64,7 +174,7 @@ export class JsonYamlChunker {
     } catch (error) {
       if (error instanceof ChunkLimitExceededError) throw error
       logger.info('Structured data chunking failed, falling back to text chunking')
-      return this.chunkAsText(content, new ChunkBudget(this.maxChunks))
+      return this.chunkAsText(content)
     }
   }
 
@@ -299,7 +409,11 @@ export class JsonYamlChunker {
     }
   }
 
-  private chunkAsText(content: string, budget: ChunkBudget, chunks: Chunk[] = []): Chunk[] {
+  private chunkAsText(
+    content: string,
+    budget: ChunkBudget = new ChunkBudget(this.maxChunks),
+    chunks: Chunk[] = []
+  ): Chunk[] {
     let currentChunk = ''
     let currentTokens = 0
     let startIndex = 0
@@ -362,9 +476,4 @@ export class JsonYamlChunker {
 
     return chunks
   }
-
-  static async chunkJsonYaml(content: string, options: ChunkerOptions = {}): Promise<Chunk[]> {
-    const chunker = new JsonYamlChunker(options)
-    return chunker.chunk(content)
-  }
 }

apps/sim/lib/file-parsers/yaml-parser.ts

@@ -10,7 +10,7 @@ import { measureYamlExpansion, type YamlExpansionLimits } from '@/lib/file-parse
  * the byte cap bounds output a sub-1 KB input can inflate to hundreds of MB;
  * the depth cap bounds the traversal's own working set.
  */
-const FILE_PARSER_YAML_LIMITS: YamlExpansionLimits = {
+export const FILE_PARSER_YAML_LIMITS: YamlExpansionLimits = {
   maxNodes: 5_000_000,
   maxSerializedBytes: 64 * 1024 * 1024,
   maxDepth: 500,

apps/sim/lib/knowledge/documents/document-processor.ts

@@ -268,13 +268,17 @@ export async function processDocument(
           mimeType.includes('json') ||
           mimeType.includes('yaml')
 
-        if (isJsonYaml && JsonYamlChunker.isStructuredData(content)) {
+        const jsonYamlChunks = isJsonYaml
+          ? await JsonYamlChunker.chunkStructured(content, {
+              chunkSize,
+              minCharactersPerChunk,
+              maxChunks: MAX_DOCUMENT_CHUNKS,
+            })
+          : null
+
+        if (jsonYamlChunks !== null) {
           logger.info('Using JSON/YAML chunker for structured data')
-          chunks = await JsonYamlChunker.chunkJsonYaml(content, {
-            chunkSize,
-            minCharactersPerChunk,
-            maxChunks: MAX_DOCUMENT_CHUNKS,
-          })
+          chunks = jsonYamlChunks
         } else if (StructuredDataChunker.isStructuredData(content, mimeType)) {
           logger.info('Using structured data chunker for spreadsheet/CSV content')
           const rowCount = metadata.totalRows ?? metadata.rowCount