fix(embeddings): honor per-model token limits and bound the JSON input path

Review round 1.

Batching used one 8,000-token constant for every model, inherited from the
knowledge-base engine this branch extracted. `batchByTokenLimit` truncates
any single text above the limit it is given, so that constant both sent
oversized input to models with a lower ceiling and silently dropped content
models with a higher one accept:

- Gemini declares 2,048, so a 3,000-token text passed through whole and the
  provider rejected it, surfacing as a 502. This also affected knowledge-base
  indexing on staging, which uses the same constant.
- Cohere declares 128,000, so anything past 8,000 was truncated for no reason.

Batch against the selected model's own `maxInputTokens` instead. Using the
per-input ceiling as the per-batch budget also keeps every individual text
within it.

The contract bounds the array arm of `input`, but a JSON-encoded array
arrives as a plain string and `normalizeInput` only expands it after
validation — so neither the 1,000-input cap nor the non-empty checks applied
to the reference-expression path the route was written to accept. `"[]"`
also reported success with no vectors. Re-check the normalized list so the
bounds hold for both shapes.

Author: mzxchandra

apps/sim/app/api/tools/embeddings/route.test.ts

@@ -129,6 +129,46 @@ describe('POST /api/tools/embeddings', () => {
     expect(mockEmbed).toHaveBeenCalledWith(['alpha', 'beta'], expect.anything())
   })
 
+  /**
+   * The contract bounds the array arm, but a JSON-encoded array reaches the
+   * route as a plain string and is only expanded after validation — so the
+   * bounds have to be re-applied to the normalized list or they hold for a
+   * native array body only.
+   */
+  describe('JSON-encoded array bounds', () => {
+    it('rejects a JSON array that exceeds the input count limit', async () => {
+      const many = JSON.stringify(Array.from({ length: 1001 }, (_, i) => `t${i}`))
+      const response = await post({ ...baseBody, input: many })
+
+      expect(response.status).toBe(400)
+      expect((await response.json()).error).toContain('cannot exceed 1000 texts')
+      expect(mockEmbed).not.toHaveBeenCalled()
+    })
+
+    it('rejects an empty JSON array instead of reporting success with no vectors', async () => {
+      const response = await post({ ...baseBody, input: '[]' })
+
+      expect(response.status).toBe(400)
+      expect((await response.json()).error).toContain('at least one text')
+      expect(mockEmbed).not.toHaveBeenCalled()
+    })
+
+    it('rejects a JSON array containing a blank entry', async () => {
+      const response = await post({ ...baseBody, input: '["ok","   "]' })
+
+      expect(response.status).toBe(400)
+      expect((await response.json()).error).toContain('entries cannot be empty')
+      expect(mockEmbed).not.toHaveBeenCalled()
+    })
+
+    it('accepts a JSON array within the bounds', async () => {
+      const response = await post({ ...baseBody, input: '["alpha","beta"]' })
+
+      expect(response.status).toBe(200)
+      expect(mockEmbed).toHaveBeenCalledWith(['alpha', 'beta'], expect.anything())
+    })
+  })
+
   it('embeds a non-JSON string as a single text', async () => {
     await post({ ...baseBody, input: 'just a sentence' })
     expect(mockEmbed).toHaveBeenCalledWith(['just a sentence'], expect.anything())

apps/sim/app/api/tools/embeddings/route.ts

@@ -3,6 +3,7 @@ import { getErrorMessage } from '@sim/utils/errors'
 import { type NextRequest, NextResponse } from 'next/server'
 import {
   embeddingsToolContract,
+  MAX_EMBEDDING_INPUTS,
   MAX_EMBEDDING_TOTAL_CHARS,
 } from '@/lib/api/contracts/tools/embeddings'
 import { getValidationErrorMessage, parseRequest, validationErrorResponse } from '@/lib/api/server'
@@ -64,6 +65,35 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
   const { provider, apiKey, model, input, taskType, dimensions } = parsed.data.body
 
   const texts = normalizeInput(input)
+
+  /**
+   * The contract bounds the array arm, but a JSON-encoded array arrives as a
+   * plain string and is only expanded here, after validation. Re-checking the
+   * normalized list is what makes the bounds hold for the reference-expression
+   * path too, rather than only for a native array body.
+   */
+  if (texts.length === 0) {
+    return NextResponse.json(
+      { success: false, error: 'input must contain at least one text' },
+      { status: 400 }
+    )
+  }
+  if (texts.length > MAX_EMBEDDING_INPUTS) {
+    return NextResponse.json(
+      {
+        success: false,
+        error: `input cannot exceed ${MAX_EMBEDDING_INPUTS} texts, received ${texts.length}`,
+      },
+      { status: 400 }
+    )
+  }
+  if (texts.some((text) => text.trim().length === 0)) {
+    return NextResponse.json(
+      { success: false, error: 'input entries cannot be empty' },
+      { status: 400 }
+    )
+  }
+
   const totalChars = texts.reduce((sum, text) => sum + text.length, 0)
   if (totalChars > MAX_EMBEDDING_TOTAL_CHARS) {
     return NextResponse.json(

apps/sim/lib/embeddings/client.test.ts

@@ -201,6 +201,40 @@ describe('embed', () => {
     expect(result.isBYOK).toBe(true)
   })
 
+  /**
+   * `batchByTokenLimit` truncates any single text above the limit it is given,
+   * so the limit has to be the selected model's own. One shared constant sent
+   * oversized input to the models with a lower ceiling and silently dropped
+   * content the models with a higher one would have accepted.
+   */
+  describe('per-model token limits', () => {
+    it("truncates against Gemini's lower ceiling rather than a shared constant", async () => {
+      fetchMock.mockResolvedValue(jsonResponse({ embeddings: [{ values: [1] }] }))
+      // ~10k tokens: over Gemini's 2048 ceiling, but under the old 8000 constant,
+      // so this used to reach the provider whole and come back a 502.
+      const long = 'word '.repeat(8000)
+
+      await embed([long], { model: 'gemini-embedding-001', apiKey: 'g-test' })
+
+      const body = JSON.parse((fetchMock.mock.calls[0][1] as RequestInit).body as string)
+      const sent = body.requests[0].content.parts[0].text
+      expect(sent.length).toBeLessThan(long.length)
+    })
+
+    it("keeps text intact up to Cohere's much higher ceiling", async () => {
+      fetchMock.mockResolvedValue(
+        jsonResponse({ embeddings: { float: [[1]] }, meta: { billed_units: { input_tokens: 9 } } })
+      )
+      // Over the old 8000 constant, well under Cohere's 128k, so it must survive.
+      const long = 'word '.repeat(8000)
+
+      await embed([long], { model: 'embed-v4.0', apiKey: 'c-test' })
+
+      const body = JSON.parse((fetchMock.mock.calls[0][1] as RequestInit).body as string)
+      expect(body.texts[0]).toBe(long)
+    })
+  })
+
   /**
    * The knowledge-base path rewrites resolved-secret plaintext back to
    * placeholders before inputs reach a provider. The block path projects
@@ -237,9 +271,7 @@ describe('embed', () => {
 
     it('estimates tokens from the projected values, not the originals', async () => {
       // Gemini omits usage, so the token count is estimated from what was sent.
-      fetchMock.mockResolvedValue(
-        jsonResponse({ embeddings: [{ values: [1, 2, 3] }] })
-      )
+      fetchMock.mockResolvedValue(jsonResponse({ embeddings: [{ values: [1, 2, 3] }] }))
 
       const result = await embed(['x'.repeat(400)], {
         model: 'gemini-embedding-001',

apps/sim/lib/embeddings/client.ts

@@ -20,7 +20,6 @@ import { batchByTokenLimit, estimateTokenCount } from '@/lib/tokenization'
 
 const logger = createLogger('EmbeddingClient')
 
-const MAX_TOKENS_PER_REQUEST = 8000
 const MAX_CONCURRENT_BATCHES = envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 50)
 const EMBEDDING_REQUEST_TIMEOUT_MS = 60_000
 
@@ -185,7 +184,15 @@ export async function embed(texts: string[], options: EmbedOptions): Promise<Emb
   const taskType = options.taskType ?? 'document'
   const provider = await resolveProvider(model, options)
 
-  const tokenBatches = batchByTokenLimit(texts, MAX_TOKENS_PER_REQUEST, model)
+  /**
+   * Batched against the selected model's own ceiling rather than one shared
+   * constant. `batchByTokenLimit` truncates any single text above the limit, so
+   * a value that is too high sends oversized input the provider rejects, and one
+   * that is too low silently drops content the provider would have accepted.
+   * Using the per-input ceiling as the per-batch budget also keeps every
+   * individual text within it.
+   */
+  const tokenBatches = batchByTokenLimit(texts, provider.info.maxInputTokens, model)
   const itemLimit = provider.adapter.maxItemsPerRequest ?? provider.info.maxItemsPerRequest
   const batches = itemLimit
     ? tokenBatches.flatMap((batch) => splitByItemLimit(batch, itemLimit))