From 0edee67a1b4cc87d01097f5bb86a3053b8cf5be4 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 13 Oct 2026 12:19:36 -0700 Subject: [PATCH] fix(knowledge): bill Gemini spend for image reads that fail the pass --- .../confluence-images.integration.ts | 96 ++++++++++--- .../documents/document-indexing-usage.test.ts | 60 +++++--- .../document-processing-source.test.ts | 2 + .../knowledge/documents/document-processor.ts | 6 +- .../documents/image-document-routing.test.ts | 16 +-- .../documents/image-extraction.test.ts | 74 ++++++---- .../knowledge/documents/image-extraction.ts | 33 +++-- apps/sim/lib/knowledge/documents/service.ts | 128 ++++++++++-------- .../documents/vision-checkpoints.test.ts | 35 ++++- .../knowledge/documents/vision-checkpoints.ts | 13 +- 10 files changed, 320 insertions(+), 143 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/confluence-images.integration.ts b/apps/sim/lib/knowledge/__integration__/confluence-images.integration.ts index 8d7280230c6..12c0c683b73 100644 --- a/apps/sim/lib/knowledge/__integration__/confluence-images.integration.ts +++ b/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(), decorative: new Set(), + cutOff: new Set(), 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()), - 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()), + 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() + return { + ...actual, + extractImageText: async (input: Parameters[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) => 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> 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', diff --git a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts index 2a713fcfc54..b76a51d43a2 100644 --- a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts +++ b/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 }> + ) { + 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,7 +442,7 @@ 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 }, }), ]) @@ -435,28 +450,35 @@ describe('knowledge image extraction usage', () => { 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() }) diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index 6703af159ac..11b71f0414b 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/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, diff --git a/apps/sim/lib/knowledge/documents/document-processor.ts b/apps/sim/lib/knowledge/documents/document-processor.ts index 081564cd243..eacda690ff2 100644 --- a/apps/sim/lib/knowledge/documents/document-processor.ts +++ b/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) { diff --git a/apps/sim/lib/knowledge/documents/image-document-routing.test.ts b/apps/sim/lib/knowledge/documents/image-document-routing.test.ts index 4c513631e36..e327c9ed11f 100644 --- a/apps/sim/lib/knowledge/documents/image-document-routing.test.ts +++ b/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' }) }) diff --git a/apps/sim/lib/knowledge/documents/image-extraction.test.ts b/apps/sim/lib/knowledge/documents/image-extraction.test.ts index f1fe235d820..f6314c2b343 100644 --- a/apps/sim/lib/knowledge/documents/image-extraction.test.ts +++ b/apps/sim/lib/knowledge/documents/image-extraction.test.ts @@ -4,16 +4,23 @@ import sharp from 'sharp' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockRequest, mockResolveKey, mockAdmission, mockCooldown, mockLoad, mockSave } = vi.hoisted( - () => ({ - mockRequest: vi.fn(), - mockResolveKey: vi.fn(), - mockAdmission: vi.fn(), - mockCooldown: vi.fn(), - mockLoad: vi.fn(), - mockSave: vi.fn(), - }) -) +const { + mockRequest, + mockResolveKey, + mockAdmission, + mockCooldown, + mockLoad, + mockSave, + mockOnUsage, +} = vi.hoisted(() => ({ + mockRequest: vi.fn(), + mockResolveKey: vi.fn(), + mockAdmission: vi.fn(), + mockCooldown: vi.fn(), + mockLoad: vi.fn(), + mockSave: vi.fn(), + mockOnUsage: vi.fn(), +})) vi.mock('@/lib/knowledge/documents/image-extraction-gemini', async (importOriginal) => ({ ...(await importOriginal()), @@ -81,10 +88,16 @@ async function extract(width = 400, height = 300, extra: Record return extractImageText({ source: await png(width, height), filename: 'Launch calendar', + onUsage: mockOnUsage, ...extra, }) } +function reportedUsage() { + expect(mockOnUsage).toHaveBeenCalledTimes(1) + return mockOnUsage.mock.calls[0][0] +} + describe('extractImageText', () => { beforeEach(() => { mockLoad.mockResolvedValue(null) @@ -98,19 +111,23 @@ describe('extractImageText', () => { it('keeps only the searchable text, with its cost, and checkpoints the result', async () => { const result = await extract(400, 300, { workspaceId: 'ws-1', checkpoint: CHECKPOINT }) - expect(result).toMatchObject({ - decorative: false, - renderedText: 'Region: EMEA. Quarter: Q1.', - usage: { model: DEFAULT_IMAGE_EXTRACTION_MODEL, inputTokens: 1200, outputTokens: 300 }, + const usage = reportedUsage() + expect(result).toMatchObject({ decorative: false, renderedText: 'Region: EMEA. Quarter: Q1.' }) + expect(usage).toMatchObject({ + model: DEFAULT_IMAGE_EXTRACTION_MODEL, + inputTokens: 1200, + outputTokens: 300, }) - expect(result.usage?.cost).toBeGreaterThan(0) + expect(usage.cost).toBeGreaterThan(0) expect(mockRequest).toHaveBeenCalledTimes(1) - expect(mockSave).toHaveBeenCalledWith(result, undefined) + // The checkpoint carries the same charge, so a resumed pass records it under the same key. + expect(mockSave).toHaveBeenCalledWith({ ...result, usage }, undefined) }) it('charges nothing when the workspace supplied its own Gemini key', async () => { mockResolveKey.mockResolvedValue({ apiKey: 'workspace-key', isBYOK: true }) - expect((await extract()).usage?.cost).toBe(0) + await extract() + expect(reportedUsage().cost).toBe(0) }) it('downscales an image wider than the model edge before sending it', async () => { @@ -123,14 +140,14 @@ describe('extractImageText', () => { it('returns a decorative image with no text but keeps its cost', async () => { respond(extraction({ kind: 'decorative', rendered_text: 'Acme logo' })) - const result = await extract(300, 120) - expect(result).toMatchObject({ decorative: true, renderedText: '' }) - expect(result.usage?.cost).toBeGreaterThan(0) + expect(await extract(300, 120)).toMatchObject({ decorative: true, renderedText: '' }) + expect(reportedUsage().cost).toBeGreaterThan(0) }) it('treats an image under the pixel floor as decorative without calling the model', async () => { expect(await extract(48, 400)).toEqual({ decorative: true, renderedText: '' }) expect(mockRequest).not.toHaveBeenCalled() + expect(mockOnUsage).not.toHaveBeenCalled() }) it('resumes from a checkpoint without a second paid call', async () => { @@ -138,6 +155,7 @@ describe('extractImageText', () => { decorative: false, renderedText: 'cached text', usage: { + chargeId: 'charge-1', model: DEFAULT_IMAGE_EXTRACTION_MODEL, cost: 0.004, inputTokens: 1, @@ -145,8 +163,12 @@ describe('extractImageText', () => { }, } mockLoad.mockResolvedValue(cached) - expect(await extract(400, 300, { checkpoint: CHECKPOINT })).toEqual(cached) + expect(await extract(400, 300, { checkpoint: CHECKPOINT })).toMatchObject({ + decorative: false, + renderedText: 'cached text', + }) expect(mockRequest).not.toHaveBeenCalled() + expect(reportedUsage()).toEqual(cached.usage) }) it('refuses bytes that are not a supported raster image', async () => { @@ -172,10 +194,8 @@ describe('extractImageText', () => { await vi.waitFor(() => expect(mockRequest).toHaveBeenCalledTimes(1)) await vi.advanceTimersByTimeAsync(10_000) /** Both replies were billed by Gemini, so the usage covers both. */ - await expect(pending).resolves.toMatchObject({ - decorative: false, - usage: { inputTokens: 2400, outputTokens: 600 }, - }) + await expect(pending).resolves.toMatchObject({ decorative: false }) + expect(reportedUsage()).toMatchObject({ inputTokens: 2400, outputTokens: 600 }) expect(mockRequest).toHaveBeenCalledTimes(2) } finally { vi.useRealTimers() @@ -191,6 +211,9 @@ describe('extractImageText', () => { for (let attempt = 0; attempt < 3; attempt++) await vi.advanceTimersByTimeAsync(10_000) expect(await pending).toMatchObject({ code: 'image_extraction_invalid' }) expect(mockRequest).toHaveBeenCalledTimes(4) + const usage = reportedUsage() + expect(usage).toMatchObject({ inputTokens: 4800, outputTokens: 1200 }) + expect(usage.cost).toBeGreaterThan(0) } finally { vi.useRealTimers() } @@ -242,6 +265,7 @@ describe('extractImageText', () => { await expect(extract()).rejects.toBe(unexpected) /** One call per failure above: none of them is retried within the extraction. */ expect(mockRequest).toHaveBeenCalledTimes(6) + expect(mockOnUsage).not.toHaveBeenCalled() }) it('retries a transient Gemini server error within the same extraction', async () => { diff --git a/apps/sim/lib/knowledge/documents/image-extraction.ts b/apps/sim/lib/knowledge/documents/image-extraction.ts index c0eed783044..9250abaeda7 100644 --- a/apps/sim/lib/knowledge/documents/image-extraction.ts +++ b/apps/sim/lib/knowledge/documents/image-extraction.ts @@ -1,5 +1,6 @@ import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' +import { generateId } from '@sim/utils/id' import { z } from 'zod' import { isTimeoutError } from '@/lib/core/execution-limits/types' import { @@ -133,11 +134,10 @@ const extractionSchema = z.object({ const { $schema: _schemaVersion, ...IMAGE_EXTRACTION_RESPONSE_SCHEMA } = z.toJSONSchema(extractionSchema) -/** What indexing keeps from one image; an image skipped before any model call carries no usage. */ +/** What indexing keeps from one image. */ interface ImageExtractionResult { decorative: boolean renderedText: string - usage?: VisionExtractionUsage } const DECORATIVE: ImageExtractionResult = { decorative: true, renderedText: '' } @@ -268,6 +268,8 @@ export async function canExtractImages(workspaceId: string | null): Promise void }): Promise { const { source, filename, workspaceId, signal } = input signal?.throwIfAborted() @@ -304,7 +307,10 @@ export async function extractImageText(input: { }) : undefined const cached = await checkpoint?.load(signal) - if (cached) return cached + if (cached) { + input.onUsage?.(cached.usage) + return cached + } const deadlineAt = input.deadlineAt ?? Number.POSITIVE_INFINITY /** Time left for the admission wait once the call and the post-call work are reserved. */ @@ -334,6 +340,7 @@ export async function extractImageText(input: { let extraction: z.infer /** Every attempt that returned a reply was billed by Gemini, including malformed ones. */ const spent = { inputTokens: 0, outputTokens: 0 } + let usage: VisionExtractionUsage try { extraction = await retryWithExponentialBackoff( async (attemptSignal) => { @@ -393,21 +400,23 @@ export async function extractImageText(input: { await recordProviderCooldown(admission, error.retryAfterMs ?? 1000) } throw classifyGeminiImageError(error, model) + } finally { + usage = { + chargeId: generateId(), + model, + cost: calculateBillableModelCost(model, spent.inputTokens, spent.outputTokens, { + isBYOK: key.isBYOK, + }).total, + ...spent, + } + if (spent.inputTokens > 0 || spent.outputTokens > 0) input.onUsage?.(usage) } const decorative = extraction.kind === 'decorative' - const { total: cost } = calculateBillableModelCost(model, spent.inputTokens, spent.outputTokens, { - isBYOK: key.isBYOK, - }) const result: VisionExtractionCheckpoint = { decorative, renderedText: decorative ? '' : extraction.rendered_text.trim(), - usage: { - model, - cost, - inputTokens: spent.inputTokens, - outputTokens: spent.outputTokens, - }, + usage, } try { await checkpoint?.save(result, signal) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 2d71d063a01..04cee6f0cea 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -11,6 +11,7 @@ import { import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' import { getErrorMessage, toError } from '@sim/utils/errors' +import { interruptibleSleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { tasks } from '@trigger.dev/sdk' import { @@ -381,6 +382,21 @@ const LARGE_DOC_CONFIG = { const HARD_DELETE_DOCUMENT_BATCH_SIZE = 250 +/** How long a timed-out or cancelled pass waits for its aborted work to report billed spend. */ +const ABORTED_PROCESSING_DRAIN_MS = 5_000 + +/** Waits for `work` to settle, but no longer than `ms`; never rejects. */ +async function settledWithin(work: Promise, ms: number): Promise { + const settled = new AbortController() + await Promise.race([ + work.then( + () => settled.abort(), + () => settled.abort() + ), + interruptibleSleep(ms, settled.signal), + ]) +} + async function withTimeout( run: (signal: AbortSignal) => Promise, timeoutMs: number, @@ -1867,7 +1883,7 @@ export async function processDocumentAsync( ) } let billableEmbeddingTokens = 0 - let visionUsage: VisionExtractionUsage | undefined + const visionCharges: Promise[] = [] let embeddingModelName = kbEmbeddingModel let embeddingPricingId = kbEmbeddingModel @@ -1884,14 +1900,61 @@ export async function processDocumentAsync( currentSourceFileProvenance ) + /** One indexing charge on the document's billing context; callers decide when to charge. */ + const recordIndexingUsage = async (entry: { + description: string + cost: number + sourceReference: string + inputTokens: number + outputTokens: number + }) => { + await recordUsage({ + userId: documentActorUserId, + workspaceId: ctx.workspaceId ?? undefined, + ...toBillingContext(billingAttribution), + entries: [ + { + category: 'model', + source: 'knowledge-base', + description: entry.description, + cost: entry.cost, + sourceReference: entry.sourceReference, + metadata: { inputTokens: entry.inputTokens, outputTokens: entry.outputTokens }, + }, + ], + }) + await checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) + } + + /** Recorded as soon as reported, since Gemini bills replies that then fail the pass. */ + const recordVisionUsage = async (visionUsage: VisionExtractionUsage) => { + if (visionUsage.cost <= 0) return + try { + await recordIndexingUsage({ + description: visionUsage.model, + cost: visionUsage.cost, + sourceReference: ['knowledge-document-vision', documentId, visionUsage.chargeId].join( + ':' + ), + inputTokens: visionUsage.inputTokens, + outputTokens: visionUsage.outputTokens, + }) + } catch (billingError) { + logger.error(`[${documentId}] Failed to record image extraction usage`, { + error: billingError, + }) + } + } + let processingCommitted = false const processingDeadlineAt = Math.min( startTime + TIMEOUTS.OVERALL_PROCESSING - 15_000, (attemptContext?.deadlineAt ?? Number.POSITIVE_INFINITY) - 15_000 ) + let processing: Promise = Promise.resolve() await withTimeout( (signal) => - runWithKnowledgeModelInputProvenance( + (processing = runWithKnowledgeModelInputProvenance( documentSecretContext.registry, async () => { await assertKnowledgeEmbeddingCapacity({ @@ -1911,6 +1974,9 @@ export async function processDocumentAsync( signal, processingDeadlineAt, imageExtraction: imageExtractionPolicyFor(ctx.connectorType), + onVisionUsage: (usage) => { + visionCharges.push(recordVisionUsage(usage)) + }, ...(indexingPassId ? { ocrCheckpoint: { knowledgeBaseId, documentId, indexingPassId } } : {}), @@ -1922,7 +1988,6 @@ export async function processDocumentAsync( signal.throwIfAborted() assertDocumentChunkCountWithinLimit(processed.chunks.length) - visionUsage = processed.metadata.visionUsage const now = new Date() @@ -2203,11 +2268,14 @@ export async function processDocumentAsync( documentSecretContext.provenance.status === 'exact' && documentSecretContext.provenance.entries.length === 0, } - ), + )), Math.max(1, processingDeadlineAt - Date.now()), 'Document processing', attemptContext?.signal - ) + ).finally(async () => { + await settledWithin(processing, ABORTED_PROCESSING_DRAIN_MS) + await Promise.all(visionCharges) + }) if (!processingCommitted) { logger.info(`[${documentId}] Discarded output from an obsolete processing attempt`) @@ -2217,32 +2285,6 @@ export async function processDocumentAsync( const processingTime = Date.now() - startTime logger.info(`[${documentId}] Successfully processed document in ${processingTime}ms`) - /** One indexing charge on the document's billing context; callers decide when to charge. */ - const recordIndexingUsage = async (entry: { - description: string - cost: number - sourceReference: string - inputTokens: number - outputTokens: number - }) => { - await recordUsage({ - userId: documentActorUserId, - workspaceId: ctx.workspaceId ?? undefined, - ...toBillingContext(billingAttribution), - entries: [ - { - category: 'model', - source: 'knowledge-base', - description: entry.description, - cost: entry.cost, - sourceReference: entry.sourceReference, - metadata: { inputTokens: entry.inputTokens, outputTokens: entry.outputTokens }, - }, - ], - }) - await checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) - } - if (billableEmbeddingTokens > 0) { try { const costMultiplier = getCostMultiplier() @@ -2300,30 +2342,6 @@ export async function processDocumentAsync( } } - if (visionUsage && visionUsage.cost > 0) { - try { - /** - * The vision call is billed whether or not the image produced chunks, so it - * sits outside the embedding gate. Its own source reference keeps the - * embedding charge's event key unchanged and collapses retries of one pass. - */ - await recordIndexingUsage({ - description: visionUsage.model, - cost: visionUsage.cost, - sourceReference: [ - 'knowledge-document-vision', - documentId, - indexingPassId ?? `model:${visionUsage.model}`, - ].join(':'), - inputTokens: visionUsage.inputTokens, - outputTokens: visionUsage.outputTokens, - }) - } catch (billingError) { - logger.error(`[${documentId}] Failed to record image extraction usage`, { - error: billingError, - }) - } - } return { outcome: 'indexed' } }) } catch (error) { diff --git a/apps/sim/lib/knowledge/documents/vision-checkpoints.test.ts b/apps/sim/lib/knowledge/documents/vision-checkpoints.test.ts index ae6194632f6..32181bbcd8a 100644 --- a/apps/sim/lib/knowledge/documents/vision-checkpoints.test.ts +++ b/apps/sim/lib/knowledge/documents/vision-checkpoints.test.ts @@ -10,6 +10,10 @@ vi.mock('@/lib/knowledge/model-input-provenance', () => ({ import type { OutboxEventContext } from '@/lib/core/outbox/service' import { VISION_CHECKPOINT_CLEANUP_EVENT } from '@/lib/knowledge/documents/checkpoint-events' +import { + readFramedCheckpoint, + writeFramedCheckpoint, +} from '@/lib/knowledge/documents/processing-checkpoint-store' import { cleanupVisionCheckpoint, createVisionCheckpoint, @@ -20,7 +24,13 @@ const source = Buffer.from('image bytes') const result = { decorative: false, renderedText: 'Region: EMEA. Quarter: Q1.', - usage: { model: 'gemini-3.8-flash', cost: 0.004, inputTokens: 1200, outputTokens: 300 }, + usage: { + chargeId: 'charge-1', + model: 'gemini-3.8-flash', + cost: 0.004, + inputTokens: 1200, + outputTokens: 300, + }, } interface CleanupRow { @@ -93,6 +103,29 @@ describe('vision extraction checkpoints', () => { expect(await checkpoint().load()).toEqual(result) }) + it('gives a checkpoint written before charge ids one stable charge, so its replays dedupe', async () => { + await checkpoint().save(result) + const [key] = [...objects.keys()] + const stored = await readFramedCheckpoint(key, { maxBodyBytes: 4 * 1024 * 1024 }) + if (!stored) throw new Error('Expected the saved checkpoint') + const { chargeId: _chargeId, ...legacyUsage } = result.usage + await writeFramedCheckpoint( + key, + stored.header, + Buffer.from(JSON.stringify({ ...result, usage: legacyUsage })), + { + cleanupEventType: VISION_CHECKPOINT_CLEANUP_EVENT, + cleanupIdPrefix: 'knowledge-vision-cleanup', + fileName: 'vision-checkpoint.json', + } + ) + + const first = await checkpoint().load() + expect(first).toMatchObject({ renderedText: result.renderedText, usage: legacyUsage }) + expect(first?.usage.chargeId).toEqual(expect.any(String)) + expect((await checkpoint().load())?.usage.chargeId).toBe(first?.usage.chargeId) + }) + it.each([ { source: Buffer.from('changed image bytes') }, { filename: 'renamed.png' }, diff --git a/apps/sim/lib/knowledge/documents/vision-checkpoints.ts b/apps/sim/lib/knowledge/documents/vision-checkpoints.ts index 661fd260876..88f62e70d87 100644 --- a/apps/sim/lib/knowledge/documents/vision-checkpoints.ts +++ b/apps/sim/lib/knowledge/documents/vision-checkpoints.ts @@ -14,6 +14,8 @@ const MAX_CHECKPOINT_BYTES = 4 * 1024 * 1024 const KEY_PATTERN = /^knowledge-vision-checkpoints\/v1\/[a-f0-9]{64}\.json$/ const visionUsageSchema = z.object({ + /** Dedupe key for one extraction's charge; a resumed pass replays it, so the charge records once. */ + chargeId: z.string().min(1), model: z.string().min(1), cost: z.number().finite().nonnegative(), inputTokens: z.number().int().nonnegative(), @@ -32,6 +34,11 @@ const checkpointBodySchema = z.object({ /** A validated extraction, reduced to what indexing keeps. */ export type VisionExtractionCheckpoint = z.infer +/** Checkpoints written before charges carried an id; their pass never billed, so one is derived. */ +const storedBodySchema = checkpointBodySchema.extend({ + usage: visionUsageSchema.extend({ chargeId: z.string().min(1).optional() }), +}) + /** * Retry cache for one image's vision extraction, never a record: a pass that dies * after the paid call but before the index swap resumes from it instead of paying @@ -71,8 +78,10 @@ export function createVisionCheckpoint(options: { }) if (!stored || stored.header.identity !== identity) return null try { - const parsed = checkpointBodySchema.safeParse(JSON.parse(stored.body.toString('utf8'))) - return parsed.success ? parsed.data : null + const parsed = storedBodySchema.safeParse(JSON.parse(stored.body.toString('utf8'))) + if (!parsed.success) return null + const { usage } = parsed.data + return { ...parsed.data, usage: { ...usage, chargeId: usage.chargeId ?? identity } } } catch { return null }