refactor(knowledge): tighten image extraction after a full review

- Every Gemini attempt takes its own admission, and 429/503 records a provider
  cooldown for the key; the retry delay is read from the body's RetryInfo when
  no Retry-After header is sent.
- The image is base64-encoded once per extraction, rotated upright even when it
  already fits, and flattened onto white before a JPEG fallback.
- Cost goes through calculateBillableModelCost; both indexing charges share one
  recordIndexingUsage helper.
- Reconciliation waits until queued image removal finishes, and the removal walks
  by id instead of rescanning from the start.
- The per-sync image cap is at least 1, so a sync can always finish.
- canExtractImages lives with the extractor; isSemanticImage treats an undeclared
  policy as OCR, so callers drop their fallbacks.
- Tests: unique embedding vectors, restored pool keys, waits for dispatched syncs,
  fake timers in finally, native noise fixtures; drops mock-only assertions.

Author: waleedlatif1

apps/docs/content/docs/platform/self-hosting/environment-variables.mdx

@@ -193,7 +193,7 @@ See [Observability](/platform/self-hosting/observability).
 | Variable | Description |
 |----------|-------------|
 | `KB_IMAGE_EXTRACTION_MODEL` | Gemini model that reads Confluence image attachments into searchable text. Defaults to `gemini-3.8-flash`; a value that is not a Gemini 3 model in the catalog falls back to the default. Needs `GEMINI_API_KEY` (or a workspace Gemini key); without one, image attachments are left out of syncs |
-| `KB_IMAGE_EXTRACTIONS_PER_SYNC` | New or changed images one connector sync sends for reading. Defaults to `500`; the rest wait for the next sync, which reports as partial |
+| `KB_IMAGE_EXTRACTIONS_PER_SYNC` | New or changed images one connector sync sends for reading, at least `1`. Defaults to `500`; the rest wait for the next sync, which reports as partial |
 | `KB_CONFIG_VISION_REQUESTS_PER_MINUTE` | Image reads per minute per Gemini key, shared by every sync on the deployment. Defaults to `300`; lower it to match a key with a smaller Gemini quota |
 | `KB_EMBEDDING_MODEL` | Embedding model for new knowledge bases. Defaults to `text-embedding-3-small`; use `ollama/<model>` for a model on your own Ollama. An unrecognised hosted model id falls back to the default, but an `ollama/` id is taken at face value — if that model is not on the server, knowledge-base creation fails rather than falling back |
 | `EMBEDDING_OUTPUT_DIMS` | Vector width new knowledge bases are stored at: `384`, `768`, `1024`, `1536` (default), or `3072`. It must be a width the chosen model can emit; anything else falls back to `1536` with a warning |

apps/sim/.env.example

@@ -259,6 +259,12 @@ CRON_SECRET=your_cron_secret # Use `openssl rand -hex 32` to generate. Authentic
 # keys in the same organization to one group using SHA-256 fingerprints, never raw keys.
 # MISTRAL_OCR_QUOTA_GROUPS={"<64-character lowercase key fingerprint>":"organization-id"}
 
+# Knowledge base images read with Gemini (Confluence "Image Attachments: Read with AI").
+# Needs GEMINI_API_KEY or a workspace Gemini key; without one, images are left out of syncs.
+# KB_IMAGE_EXTRACTION_MODEL=gemini-3.8-flash
+# KB_IMAGE_EXTRACTIONS_PER_SYNC=500
+# KB_CONFIG_VISION_REQUESTS_PER_MINUTE=300
+
 # Official Sim Search Slack app (optional; requires existing Search access)
 # Supply these through the deployment environment (for example, ECS Secrets Manager injection).
 # SLACK_SEARCH_APP_ID=

apps/sim/app/workspace/[workspaceId]/knowledge/[id]/[documentId]/document.tsx

@@ -402,7 +402,7 @@ export function Document({
   const DocumentIcon =
     ConnectorIcon || getDocumentIcon(documentData?.mimeType ?? '', effectiveDocumentName)
   const isCompleted = documentData?.processingStatus === 'completed'
-  /** Only a decorative image (a logo or icon) completes with nothing indexed. */
+  /** A completed image with no chunks held nothing searchable, such as a logo or an icon. */
   const isImageWithNothingToIndex =
     isCompleted && documentData.mimeType.startsWith('image/') && documentData.chunkCount === 0
 

apps/sim/connectors/confluence/attachments.ts

@@ -430,8 +430,9 @@ export async function listConfluenceAttachments(
     state.serverFailureStreak = 0
     const { page } = listing
     for (const attachment of page.results) {
-      if (attachment.status !== 'current' || !attachmentMimeType(attachment)) continue
-      if (isImageAttachment(attachment) && (!input.includeImages || isBelowImageFloor(attachment)))
+      const mimeType = attachmentMimeType(attachment)
+      if (attachment.status !== 'current' || !mimeType) continue
+      if (mimeType.startsWith('image/') && (!input.includeImages || isBelowImageFloor(attachment)))
         continue
       if (!sameParent(attachmentParent(attachment), parent)) {
         throw new Error('Confluence returned an attachment belonging to another parent')
@@ -595,7 +596,7 @@ export async function getConfluenceAttachment(
     }
   }
   if (
-    isImageAttachment(attachment) &&
+    mimeType.startsWith('image/') &&
     !isImageExtractionEnabled(confluenceConnectorMeta.imageExtraction, sourceConfig)
   ) {
     return {

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

@@ -28,6 +28,8 @@ const fixture = vi.hoisted(() => {
     images: new Map<string, number>(),
     decorative: new Set<string>(),
     extractions: [] as string[],
+    /** Each chunk gets its own one-hot vector, so the shared HNSW graph sees no duplicates. */
+    nextVector: 0,
   }
 })
 vi.mock('@/lib/uploads/core/setup.server', () => ({
@@ -39,15 +41,20 @@ vi.mock('@/lib/embeddings', async () => ({
   ...(await import('@/lib/embeddings/client')),
   assertKnowledgeEmbeddingCapacity: async () => {},
   embedKnowledge: async (texts: string[]) => ({
-    embeddings: texts.map(() => [1, ...Array<number>(1535).fill(0)]),
+    embeddings: texts.map(() => {
+      const vector = Array<number>(1536).fill(0)
+      vector[fixture.nextVector++ % 512] = 1
+      return vector
+    }),
     totalTokens: texts.length,
     billableTokens: 0,
     isBYOK: true,
     modelName: 'text-embedding-3-small',
     pricingId: 'text-embedding-3-small',
   }),
 }))
-vi.mock('@/lib/knowledge/documents/image-extraction', () => ({
+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$/, '')
@@ -117,7 +124,12 @@ import { IMAGES_DEFERRED_NOTICE } from '@/lib/knowledge/connectors/sync-limits'
 describe('Confluence image attachments through sync, indexing and reconciliation', () => {
   const ids = createKnowledgeAclFixtureIds()
   let billing: Awaited<ReturnType<typeof resolveBillingAttribution>>
-  const priorGeminiKey = env.GEMINI_API_KEY
+  const priorGeminiKeys = {
+    GEMINI_API_KEY: env.GEMINI_API_KEY,
+    GEMINI_API_KEY_1: env.GEMINI_API_KEY_1,
+    GEMINI_API_KEY_2: env.GEMINI_API_KEY_2,
+    GEMINI_API_KEY_3: env.GEMINI_API_KEY_3,
+  }
 
   async function configure(imageAttachments: 'semantic' | 'off' | undefined) {
     await db
@@ -135,34 +147,38 @@ describe('Confluence image attachments through sync, indexing and reconciliation
       })
       .where(eq(knowledgeConnector.id, ids.connectorId))
   }
-  /** 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' },
-      input: {
-        connectorId: ids.connectorId,
-        updates: {
-          sourceConfig: { domain: 'fixture.atlassian.net', spaceKey: 'ENG', imageAttachments },
-        },
-      },
-    })
+  /** Waits for the sync a create or update dispatched in-process to land. */
+  async function syncFinished(connectorId: string) {
     await expect
       .poll(
         async () => {
           const [row] = await db
             .select()
             .from(knowledgeConnector)
-            .where(eq(knowledgeConnector.id, ids.connectorId))
+            .where(eq(knowledgeConnector.id, connectorId))
           return {
             status: row?.status,
             error: row?.lastSyncError,
             synced: Boolean(row?.lastSyncAt),
           }
         },
-        { timeout: 30_000 }
+        { timeout: 45_000 }
       )
       .toEqual({ status: 'active', error: null, synced: true })
   }
+  /** 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' },
+      input: {
+        connectorId: ids.connectorId,
+        updates: {
+          sourceConfig: { domain: 'fixture.atlassian.net', spaceKey: 'ENG', imageAttachments },
+        },
+      },
+    })
+    await syncFinished(ids.connectorId)
+  }
   async function imageRemovalPending() {
     const [row] = await db
       .select({ pending: knowledgeConnector.imageRemovalPending })
@@ -221,7 +237,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     await configure('semantic')
   })
   afterAll(async () => {
-    Object.assign(env, { GEMINI_API_KEY: priorGeminiKey })
+    Object.assign(env, priorGeminiKeys)
     await deleteWorkspaceFixture(db, eq(workspace.id, ids.workspaceId))
     await db.delete(organization).where(eq(organization.id, ids.organizationId))
     await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId]))
@@ -259,7 +275,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     await sync()
     expect((await images()).every((row) => row.contentHash?.endsWith(':2'))).toBe(true)
     expect(await lastSyncError()).toBeNull()
-  })
+  }, 120_000)
 
   it('does not read unchanged images again on a full resync', async () => {
     fixture.images.set('image-a', 1)
@@ -288,7 +304,7 @@ describe('Confluence image attachments through sync, indexing and reconciliation
     fixture.extractions.length = 0
     await setImageAttachments('semantic')
     expect(fixture.extractions.sort()).toEqual(['image-a.png', 'image-b.png'])
-  })
+  }, 120_000)
 
   it('keeps images off for a connector created before the setting existed', async () => {
     fixture.images.set('image-a', 1)
@@ -347,7 +363,9 @@ describe('Confluence image attachments through sync, indexing and reconciliation
         .from(knowledgeConnector)
         .where(eq(knowledgeConnector.id, created.connector.id))
       expect(row?.sourceConfig).toMatchObject({ imageAttachments: saved })
+      await syncFinished(created.connector.id)
       await db.delete(knowledgeConnector).where(eq(knowledgeConnector.id, created.connector.id))
-    }
+    },
+    120_000
   )
 })

apps/sim/lib/knowledge/connectors/image-admission.ts

@@ -1,11 +1,10 @@
 import { env, envNumber } from '@/lib/core/config/env'
-import { resolveGeminiImageKey } from '@/lib/knowledge/documents/image-extraction-gemini'
 import type { ImageExtractionPolicy } from '@/lib/knowledge/documents/image-extraction-policy'
 import { getConnectorMeta } from '@/connectors/registry'
 
 /** Images one sync may send for vision extraction, new or changed; the rest wait for the next sync. */
 const IMAGE_EXTRACTIONS_PER_SYNC = envNumber(env.KB_IMAGE_EXTRACTIONS_PER_SYNC, 500, {
-  min: 0,
+  min: 1,
   integer: true,
 })
 
@@ -14,14 +13,6 @@ export function imageExtractionPolicyFor(connectorType: string | null): ImageExt
   return connectorType && getConnectorMeta(connectorType)?.imageExtraction ? 'semantic' : 'ocr'
 }
 
-/**
- * Whether the workspace can have images read at all: its own Gemini key, or one
- * on the deployment. Without either, images are left unlisted rather than failed.
- */
-export async function canExtractImages(workspaceId: string | null): Promise<boolean> {
-  return (await resolveGeminiImageKey(workspaceId)) !== null
-}
-
 /**
  * The part of a listing checkpoint that counts this sync's image extractions.
  * Living on the checkpoint keeps the count across a sync that yields and resumes

apps/sim/lib/knowledge/connectors/member-sync-engine.ts

@@ -1754,7 +1754,7 @@ async function syncDedicatedMemberContent(input: {
   })
   return {
     complete: pass.complete,
-    imageExtractionsAdmitted: pass.checkpoint.imageExtractionsAdmitted ?? 0,
+    imageExtractionsAdmitted: pass.checkpoint.imageExtractionsAdmitted,
   }
 }
 

apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts

@@ -94,7 +94,7 @@ interface StoredPage {
   sourceSeenAt: Date | null
 }
 
-/** Images on, as for a new connector, so only the image opt-out test runs its removal. */
+/** Images on, as a new connector is created. */
 const SOURCE_CONFIG = {
   domain: 'fixture.atlassian.net',
   spaceKey: 'ENG',

apps/sim/lib/knowledge/connectors/sync-content-pass.ts

@@ -1,6 +1,6 @@
 import { db } from '@sim/db'
 import { document, knowledgeConnector } from '@sim/db/schema'
-import { and, eq, inArray, isNotNull, isNull, lt, type SQL, sql } from 'drizzle-orm'
+import { and, eq, gt, inArray, isNotNull, isNull, lt, type SQL, sql } from 'drizzle-orm'
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import type { DbOrTx } from '@/lib/db/types'
 import {
@@ -12,7 +12,10 @@ import {
   GOOGLE_COMPANY_PERMISSION_REFRESH_MS,
   isGoogleCompanySource,
 } from '@/lib/knowledge/connectors/google-company-scheduler'
-import { imageExtractionPolicyFor } from '@/lib/knowledge/connectors/image-admission'
+import {
+  type ImageAdmissionState,
+  imageExtractionPolicyFor,
+} from '@/lib/knowledge/connectors/image-admission'
 import {
   beginListingCheckpoint,
   type ListingCheckpoint,
@@ -98,6 +101,14 @@ interface ContentPassInput {
   ) => Promise<{ permissionsIncomplete: boolean } | undefined>
 }
 
+/** The image fields of a listing checkpoint that a page persists alongside its seen stamps. */
+function imageCountOf(state: ImageAdmissionState): ImageAdmissionState {
+  return {
+    imageExtractionsAdmitted: state.imageExtractionsAdmitted,
+    imagesDeferred: state.imagesDeferred,
+  }
+}
+
 /** One durable content cycle shared by content-owned and member-visibility connectors. */
 export async function runConnectorContentPass(input: ContentPassInput) {
   const { matchContentHash } = input.connectorConfig
@@ -143,10 +154,7 @@ export async function runConnectorContentPass(input: ContentPassInput) {
     })
   }
   let hydratedCount = 0
-  let storedImageCount = {
-    imageExtractionsAdmitted: checkpoint.imageExtractionsAdmitted,
-    imagesDeferred: checkpoint.imagesDeferred,
-  }
+  let storedImageCount = imageCountOf(checkpoint)
   const syncIntervalMinutes = effectiveConnectorSyncIntervalMinutes(
     input.documentAccess,
     input.connector.syncIntervalMinutes ?? SIM_SEARCH_SYNC_INTERVAL_MINUTES
@@ -267,10 +275,7 @@ export async function runConnectorContentPass(input: ContentPassInput) {
                 )
               )
           }
-          const imageCount = {
-            imageExtractionsAdmitted: cycle.imageExtractionsAdmitted,
-            imagesDeferred: cycle.imagesDeferred,
-          }
+          const imageCount = imageCountOf(cycle)
           if (
             imageCount.imageExtractionsAdmitted !== storedImageCount.imageExtractionsAdmitted ||
             imageCount.imagesDeferred !== storedImageCount.imagesDeferred
@@ -370,9 +375,11 @@ export async function runConnectorContentPass(input: ContentPassInput) {
     },
   })
   const imagesRemoved = await removeDisabledImages(input)
-  const reconciliation = checkpoint.complete
-    ? await reconcileCompletedListing(input, checkpoint, withLease, withAclPage)
-    : { finished: false, notice: null }
+  /** Images the removal has not reached yet are unlisted, so reconciling now would count them as absent. */
+  const reconciliation =
+    checkpoint.complete && imagesRemoved
+      ? await reconcileCompletedListing(input, checkpoint, withLease, withAclPage)
+      : { finished: false, notice: null }
   /** Unverified permissions, an incomplete listing and unrefreshed content are independent holds; an admin needs each, one per line. */
   const holdNotice =
     [
@@ -405,6 +412,8 @@ export async function removeDisabledImages(
   >
 ): Promise<boolean> {
   if (!input.connector.imageRemovalPending) return true
+  /** Walks by id through `doc_connector_reconciliation_v2_idx`, so each page resumes where the last ended. */
+  let afterId: string | undefined
   while (true) {
     if (Date.now() >= input.deadlineAt) return false
     await input.lease.beatIfDue()
@@ -416,11 +425,14 @@ export async function removeDisabledImages(
           eq(document.connectorId, input.connectorId),
           inArray(document.mimeType, [...SEMANTIC_IMAGE_MIME_TYPES]),
           eq(document.userExcluded, false),
-          isNull(document.archivedAt)
+          isNull(document.archivedAt),
+          afterId ? gt(document.id, afterId) : undefined
         )
       )
+      .orderBy(document.id)
       .limit(25)
     if (rows.length === 0) break
+    afterId = rows.at(-1)?.id
     const removed = await hardDeleteDocuments(
       rows.map((row) => row.id),
       input.runId,

apps/sim/lib/knowledge/connectors/sync-engine.ts

@@ -343,7 +343,6 @@ export interface ContentPassOutcome {
     contentFailures?: boolean
     permissionFailures?: boolean
     imagesDeferred?: boolean
-    imageExtractionsAdmitted?: number
     listingFailures?: ExternalListingFailures | null
     startedAt: string
     listedCount: number