fix(knowledge): bill Gemini spend for image reads that fail the pass (#8967)

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/confluence-images.integration.ts

@@ -1,7 +1,8 @@
 /**
  * Real sync workers, PostgreSQL, file storage, indexing and reconciliation for Confluence image
  * attachments. The source listing is a synthetic connector behind the real Confluence meta, and
- * the vision extraction and embeddings are synthetic, so no request leaves the process.
+ * the vision extraction and embeddings are synthetic, so no request leaves the process. A
+ * `cutOff` image runs the real extractor against a Gemini reply stopped at its token limit.
  */
 import { mkdtempSync } from 'node:fs'
 import { rm } from 'node:fs/promises'
@@ -13,10 +14,12 @@ import {
   embedding,
   knowledgeConnector,
   organization,
+  usageLog,
   user,
   workspace,
 } from '@sim/db/schema'
 import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
+import { createSessionPrincipal } from '@sim/testing/factories/principal.factory'
 import { and, eq, inArray, isNull } from 'drizzle-orm'
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -27,6 +30,7 @@ const fixture = vi.hoisted(() => {
     /** externalId to version; a version bump changes the image's bytes and hash. */
     images: new Map<string, number>(),
     decorative: new Set<string>(),
+    cutOff: new Set<string>(),
     extractions: [] as string[],
     /** Each chunk gets its own one-hot vector, so the shared HNSW graph sees no duplicates. */
     nextVector: 0,
@@ -53,22 +57,49 @@ vi.mock('@/lib/embeddings', async () => ({
     pricingId: 'text-embedding-3-small',
   }),
 }))
-vi.mock('@/lib/knowledge/documents/image-extraction', async (importOriginal) => ({
-  ...(await importOriginal<typeof import('@/lib/knowledge/documents/image-extraction')>()),
-  extractImageText: async ({ filename }: { filename: string }) => {
-    fixture.extractions.push(filename)
-    const externalId = filename.replace(/\.png$/, '')
-    if (fixture.decorative.has(externalId)) return { decorative: true, renderedText: '' }
-    return {
-      decorative: false,
-      renderedText: `${externalId} version ${fixture.images.get(externalId)}: Region EMEA launches in Q1.`,
-      usage: { model: 'gemini-3.8-flash', cost: 0, inputTokens: 10, outputTokens: 10 },
-    }
-  },
+vi.mock('@/lib/knowledge/documents/image-extraction-gemini', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@/lib/knowledge/documents/image-extraction-gemini')>()),
+  requestGeminiImageExtraction: async () => ({
+    text: '{"kind": "table", "rendered_text": "Region: EMEA',
+    finishReason: 'MAX_TOKENS',
+    inputTokens: 1200,
+    outputTokens: 32_768,
+  }),
 }))
+vi.mock('@/lib/knowledge/documents/image-extraction', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@/lib/knowledge/documents/image-extraction')>()
+  return {
+    ...actual,
+    extractImageText: async (input: Parameters<typeof actual.extractImageText>[0]) => {
+      const externalId = input.filename.replace(/\.png$/, '')
+      if (fixture.cutOff.has(externalId)) return actual.extractImageText(input)
+      fixture.extractions.push(input.filename)
+      if (fixture.decorative.has(externalId)) return { decorative: true, renderedText: '' }
+      return {
+        decorative: false,
+        renderedText: `${externalId} version ${fixture.images.get(externalId)}: Region EMEA launches in Q1.`,
+      }
+    },
+  }
+})
 vi.mock('@/connectors/registry.server', async () => {
   const { confluenceConnectorMeta } = await import('@/connectors/confluence/meta')
   const enabled = (config: Record<string, unknown>) => config.imageAttachments === 'semantic'
+  /** A real, noisy PNG over the semantic byte floor, so the real extractor reads it. */
+  const cutOffImage = async () => {
+    const sharp = (await import('sharp')).default
+    return sharp({
+      create: {
+        width: 256,
+        height: 256,
+        channels: 3,
+        background: '#808080',
+        noise: { type: 'gaussian', mean: 128, sigma: 60 },
+      },
+    })
+      .png()
+      .toBuffer()
+  }
   const stub = (externalId: string) => {
     const version = fixture.images.get(externalId) ?? 1
     return {
@@ -95,7 +126,9 @@ vi.mock('@/connectors/registry.server', async () => {
           externalId: string
         ) => {
           if (!fixture.images.has(externalId)) return null
-          const bytes = Buffer.from(`fixture png ${externalId} v${fixture.images.get(externalId)}`)
+          const bytes = fixture.cutOff.has(externalId)
+            ? await cutOffImage()
+            : Buffer.from(`fixture png ${externalId} v${fixture.images.get(externalId)}`)
           return {
             ...stub(externalId),
             contentDeferred: false,
@@ -121,9 +154,11 @@ import {
 import * as connectorTokens from '@/lib/knowledge/connectors/access-token'
 import { executeSync } from '@/lib/knowledge/connectors/sync-engine'
 import { IMAGES_DEFERRED_NOTICE } from '@/lib/knowledge/connectors/sync-limits'
+import { resolveImageExtractionModel } from '@/lib/knowledge/documents/image-extraction-gemini'
 
 describe('Confluence image attachments through sync, indexing and reconciliation', () => {
   const ids = createKnowledgeAclFixtureIds()
+  const principal = createSessionPrincipal({ userId: ids.aliceId, sessionId: 'confluence-images' })
   let billing: Awaited<ReturnType<typeof resolveBillingAttribution>>
   const priorGeminiKeys = {
     GEMINI_API_KEY: env.GEMINI_API_KEY,
@@ -170,7 +205,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
   /** Changes the setting the way a user does; the update dispatches the next sync, which runs in-process. */
   async function setImageAttachments(imageAttachments: 'semantic' | 'off') {
     await updateKnowledgeConnector.execute({
-      principal: { kind: 'session', userId: ids.aliceId, sessionId: 'confluence-images' },
+      principal,
       input: {
         connectorId: ids.connectorId,
         updates: {
@@ -182,7 +217,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
   }
   function setExcluded(documentId: string, operation: 'exclude' | 'restore') {
     return updateKnowledgeConnectorDocuments.execute({
-      principal: { kind: 'session', userId: ids.aliceId, sessionId: 'confluence-images' },
+      principal,
       input: {
         knowledgeBaseId: ids.knowledgeBaseId,
         connectorId: ids.connectorId,
@@ -244,6 +279,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     Object.assign(env, { GEMINI_API_KEY: 'fixture-gemini-key' })
     fixture.extractions.length = 0
     fixture.decorative.clear()
+    fixture.cutOff.clear()
     fixture.images.clear()
     await db.delete(document).where(eq(document.connectorId, ids.connectorId))
     await configure('semantic')
@@ -335,6 +371,32 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     expect(await setExcluded(excluded.id, 'restore')).toMatchObject({ count: 1 })
   }, 120_000)
 
+  it('bills the Gemini call once when a reply cut off at its token limit fails the image', async () => {
+    fixture.images.set('image-a', 1)
+    fixture.cutOff.add('image-a')
+    await executeSync(ids.connectorId, { fullSync: true, billingAttribution: billing })
+
+    await expect
+      .poll(async () => (await images()).map((row) => row.processingStatus), { timeout: 45_000 })
+      .toEqual(['failed'])
+    const [failed] = await images()
+    expect(failed.processingError).toContain('more text than one extraction can return')
+    expect(await vectors([failed.id])).toEqual([])
+    const charges = await db
+      .select()
+      .from(usageLog)
+      .where(
+        and(
+          eq(usageLog.workspaceId, ids.workspaceId),
+          eq(usageLog.source, 'knowledge-base'),
+          eq(usageLog.description, resolveImageExtractionModel())
+        )
+      )
+    expect(charges).toHaveLength(1)
+    expect(Number(charges[0].cost)).toBeGreaterThan(0)
+    expect(charges[0].metadata).toEqual({ inputTokens: 1200, outputTokens: 32_768 })
+  }, 120_000)
+
   it('keeps images off for a connector created before the setting existed', async () => {
     fixture.images.set('image-a', 1)
     await configure(undefined)
@@ -378,7 +440,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     'saves a new connector with Image Attachments on unless the request chose otherwise (%o)',
     async (chosen, saved) => {
       const created = await createKnowledgeConnector.execute({
-        principal: { kind: 'session', userId: ids.aliceId, sessionId: 'confluence-images' },
+        principal,
         input: {
           knowledgeBaseId: ids.knowledgeBaseId,
           connectorType: 'confluence',

apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts

@@ -367,6 +367,7 @@ describe('knowledge document indexing usage', () => {
 
 describe('knowledge image extraction usage', () => {
   const VISION_USAGE = {
+    chargeId: 'charge-1',
     model: 'gemini-3.8-flash',
     cost: 0.004,
     inputTokens: 1200,
@@ -411,12 +412,26 @@ describe('knowledge image extraction usage', () => {
     )
   }
 
+  /** Processing that reports a vision charge the way the extractor does, then settles. */
+  function extractThen(
+    usage: typeof VISION_USAGE,
+    settle: () => Promise<{ chunks: unknown[]; metadata: Record<string, number> }>
+  ) {
+    mockProcessDocument.mockImplementation(async (...args: unknown[]) => {
+      const access = args[6] as { onVisionUsage: (reported: typeof VISION_USAGE) => void }
+      access.onVisionUsage(usage)
+      return settle()
+    })
+  }
+
+  const NOTHING_INDEXED = {
+    chunks: [],
+    metadata: { chunkCount: 0, tokenCount: 0, characterCount: 0 },
+  }
+
   it('bills a decorative image that indexed nothing', async () => {
     armImageReads('confluence')
-    mockProcessDocument.mockResolvedValue({
-      chunks: [],
-      metadata: { chunkCount: 0, tokenCount: 0, characterCount: 0, visionUsage: VISION_USAGE },
-    })
+    extractThen(VISION_USAGE, async () => NOTHING_INDEXED)
     await process()
 
     expect(mockGenerateEmbeddings).not.toHaveBeenCalled()
@@ -427,36 +442,43 @@ describe('knowledge image extraction usage', () => {
         source: 'knowledge-base',
         description: 'gemini-3.8-flash',
         cost: 0.004,
-        sourceReference: `knowledge-document-vision:${DOCUMENT_ID}:pass-1`,
+        sourceReference: `knowledge-document-vision:${DOCUMENT_ID}:charge-1`,
         metadata: { inputTokens: 1200, outputTokens: 300 },
       }),
     ])
   })
 
   it('bills the vision call beside the embedding charge without changing it', async () => {
     armImageReads('confluence')
-    mockProcessDocument.mockResolvedValue({
+    extractThen(VISION_USAGE, async () => ({
       chunks: [{ text: 'Region: EMEA', metadata: { startIndex: 0, endIndex: 12 } }],
-      metadata: { chunkCount: 1, tokenCount: 4, characterCount: 12, visionUsage: VISION_USAGE },
-    })
+      metadata: { chunkCount: 1, tokenCount: 4, characterCount: 12 },
+    }))
     await process()
 
     expect(mockRecordUsage).toHaveBeenCalledTimes(2)
-    expect(recordedSourceReference(0)).toBe(`knowledge-document:${DOCUMENT_ID}:pass-1`)
-    expect(recordedSourceReference(1)).toBe(`knowledge-document-vision:${DOCUMENT_ID}:pass-1`)
+    expect([recordedSourceReference(0), recordedSourceReference(1)].sort()).toEqual([
+      `knowledge-document-vision:${DOCUMENT_ID}:charge-1`,
+      `knowledge-document:${DOCUMENT_ID}:pass-1`,
+    ])
+  })
+
+  it('bills the vision call when a later stage fails the pass', async () => {
+    armImageReads('confluence')
+    extractThen(VISION_USAGE, async () => ({
+      chunks: [{ text: 'Region: EMEA', metadata: { startIndex: 0, endIndex: 12 } }],
+      metadata: { chunkCount: 1, tokenCount: 4, characterCount: 12 },
+    }))
+    mockGenerateEmbeddings.mockRejectedValue(new Error('embedding provider down'))
+    await expect(process()).rejects.toThrow()
+
+    expect(mockRecordUsage).toHaveBeenCalledTimes(1)
+    expect(recordedSourceReference(0)).toBe(`knowledge-document-vision:${DOCUMENT_ID}:charge-1`)
   })
 
   it('records no vision charge when the call cost nothing, as with a workspace key', async () => {
     armImageReads('confluence')
-    mockProcessDocument.mockResolvedValue({
-      chunks: [],
-      metadata: {
-        chunkCount: 0,
-        tokenCount: 0,
-        characterCount: 0,
-        visionUsage: { ...VISION_USAGE, cost: 0 },
-      },
-    })
+    extractThen({ ...VISION_USAGE, cost: 0 }, async () => NOTHING_INDEXED)
     await process()
     expect(mockRecordUsage).not.toHaveBeenCalled()
   })

apps/sim/lib/knowledge/documents/document-processing-source.test.ts

@@ -293,6 +293,7 @@ describe('knowledge document processing source', () => {
         signal: expect.any(AbortSignal),
         processingDeadlineAt: expect.any(Number),
         imageExtraction: 'ocr',
+        onVisionUsage: expect.any(Function),
       },
       PERSISTED_CONTEXT.workspaceId,
       undefined,
@@ -336,6 +337,7 @@ describe('knowledge document processing source', () => {
         signal: expect.any(AbortSignal),
         processingDeadlineAt: expect.any(Number),
         imageExtraction: 'ocr',
+        onVisionUsage: expect.any(Function),
       },
       PERSISTED_CONTEXT.workspaceId,
       undefined,

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

@@ -236,6 +236,8 @@ export type SourceFileAccess = Pick<
   processingDeadlineAt?: number
   /** The document's source policy for images; absent keeps the OCR path. */
   imageExtraction?: ImageExtractionPolicy
+  /** Receives what the vision model billed, even when processing then fails. */
+  onVisionUsage?: (usage: VisionExtractionUsage) => void
 }
 
 type ProcessingMethod = 'file-parser' | 'mistral-ocr' | 'vision-extract'
@@ -262,7 +264,6 @@ export async function processDocument(
     characterCount: number
     processingMethod: ProcessingMethod
     cloudUrl?: string
-    visionUsage?: VisionExtractionUsage
   }
 }> {
   logger.info('Processing document', { mimeType })
@@ -277,6 +278,7 @@ export async function processDocument(
           checkpoint: access.ocrCheckpoint,
           deadlineAt: access.processingDeadlineAt,
           signal: access.signal,
+          onUsage: access.onVisionUsage,
         })
       : undefined
     if (image?.decorative) {
@@ -291,7 +293,6 @@ export async function processDocument(
           tokenCount: 0,
           characterCount: 0,
           processingMethod: 'vision-extract',
-          visionUsage: image.usage,
         },
       }
     }
@@ -413,7 +414,6 @@ export async function processDocument(
         characterCount,
         processingMethod,
         cloudUrl,
-        visionUsage: image?.usage,
       },
     }
   } catch (error) {

apps/sim/lib/knowledge/documents/image-document-routing.test.ts

@@ -43,7 +43,7 @@ import { runWithKnowledgeModelInputProvenance } from '@/lib/knowledge/model-inpu
 
 const IMAGE_URL = 'https://example.com/schedule.png'
 const IMAGE_BYTES = Buffer.from('png bytes')
-const USAGE = { model: 'gemini-3.8-flash', cost: 0.004, inputTokens: 1200, outputTokens: 300 }
+const onVisionUsage = vi.fn()
 
 function process(policy: ImageExtractionPolicy | undefined, mimeType = 'image/png') {
   return runWithKnowledgeModelInputProvenance(
@@ -53,6 +53,7 @@ function process(policy: ImageExtractionPolicy | undefined, mimeType = 'image/pn
         userId: 'user-1',
         imageExtraction: policy,
         processingDeadlineAt: Date.now() + 600_000,
+        onVisionUsage,
         ocrCheckpoint: { knowledgeBaseId: 'kb-1', documentId: 'doc-1', indexingPassId: 'pass-1' },
       }),
     { opaqueInputSafe: true }
@@ -68,7 +69,6 @@ describe('image document routing', () => {
     mockExtractImageText.mockResolvedValue({
       decorative: false,
       renderedText: 'Region: EMEA. Quarter: Q1. Launch date: March 3.',
-      usage: USAGE,
     })
   })
 
@@ -80,25 +80,23 @@ describe('image document routing', () => {
         source: IMAGE_BYTES,
         filename: 'Schedule',
         checkpoint: { knowledgeBaseId: 'kb-1', documentId: 'doc-1', indexingPassId: 'pass-1' },
+        onUsage: onVisionUsage,
       })
     )
     expect(mockExecuteMistralParse).not.toHaveBeenCalled()
     expect(processed.chunks.map((chunk) => chunk.text).join(' ')).toContain('Region: EMEA')
-    expect(processed.metadata).toMatchObject({
-      processingMethod: 'vision-extract',
-      visionUsage: USAGE,
-    })
+    expect(processed.metadata).toMatchObject({ processingMethod: 'vision-extract' })
   })
 
   it('completes a decorative image with no chunks instead of failing it as unreadable', async () => {
-    mockExtractImageText.mockResolvedValue({ decorative: true, renderedText: '', usage: USAGE })
+    mockExtractImageText.mockResolvedValue({ decorative: true, renderedText: '' })
     const processed = await process('semantic')
     expect(processed.chunks).toEqual([])
-    expect(processed.metadata).toMatchObject({ chunkCount: 0, visionUsage: USAGE })
+    expect(processed.metadata).toMatchObject({ chunkCount: 0 })
   })
 
   it('fails an informative image whose extraction came back empty', async () => {
-    mockExtractImageText.mockResolvedValue({ decorative: false, renderedText: '', usage: USAGE })
+    mockExtractImageText.mockResolvedValue({ decorative: false, renderedText: '' })
     await expect(process('semantic')).rejects.toMatchObject({ code: 'no_extractable_text' })
   })