diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index b1869c719f9..d7836ac96b6 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -278,6 +278,7 @@ jobs: lib/knowledge/__integration__/slack-empty-threads.integration.ts lib/knowledge/__integration__/kb-block-search.integration.ts lib/knowledge/__integration__/unfilled-projection-source.integration.ts + lib/knowledge/__integration__/purged-detach-reservation.integration.ts lib/core/outbox/service.integration.ts lib/knowledge/__integration__/connector-upload.integration.ts lib/uploads/contexts/organization-logo/application.integration.ts diff --git a/apps/sim/background/cleanup-soft-deletes.test.ts b/apps/sim/background/cleanup-soft-deletes.test.ts index 625a920ff47..2b0c770b70a 100644 --- a/apps/sim/background/cleanup-soft-deletes.test.ts +++ b/apps/sim/background/cleanup-soft-deletes.test.ts @@ -26,6 +26,7 @@ const { mockPrepareChatCleanup, mockResolveStorageBillingContext, mockSelectRowsByIdChunks, + mockSettleDetachedConnectorReservations, mockDeduplicateWorkflowName, mockAllocateUniqueWorkspaceFileName, mockDeduplicateFolderName, @@ -46,6 +47,7 @@ const { mockPrepareChatCleanup: vi.fn(async () => ({ execute: vi.fn(async () => undefined) })), mockResolveStorageBillingContext: vi.fn(), mockSelectRowsByIdChunks: vi.fn(async () => [] as unknown[]), + mockSettleDetachedConnectorReservations: vi.fn(async () => undefined), })) vi.mock('@/lib/billing/cleanup-dispatcher', () => ({ runCleanupWithLimits: vi.fn() })) @@ -71,6 +73,10 @@ vi.mock('@/lib/billing/storage', () => ({ resolveStorageBillingContext: mockResolveStorageBillingContext, })) +vi.mock('@/lib/knowledge/connectors/detachment', () => ({ + settleDetachedConnectorReservations: mockSettleDetachedConnectorReservations, +})) + vi.mock('@/lib/knowledge/documents/service', () => ({ hardDeleteDocuments: mockHardDeleteDocuments, })) @@ -303,6 +309,32 @@ describe('cleanup soft deletes', () => { ) }) + it('settles overdrawn reservations before the documents and the rest before the base delete', async () => { + mockChunkedBatchDelete.mockImplementationOnce( + async (options: { onBatch?: (rows: Array<{ id: string }>) => Promise }) => { + await options.onBatch?.([{ id: 'kb-1' }, { id: 'kb-2' }]) + mockKnowledgeBaseContainerDelete() + return { deleted: 2, failed: 0 } + } + ) + dbChainMockFns.limit + .mockResolvedValueOnce([{ id: 'doc-1' }]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + + await runCleanupSoftDeletes(basePayload) + + expect(mockSettleDetachedConnectorReservations.mock.calls).toEqual([ + [['kb-1', 'kb-2'], 'overdrawn'], + [['kb-1', 'kb-2'], 'remaining'], + ]) + const [overdrawn, remaining] = mockSettleDetachedConnectorReservations.mock.invocationCallOrder + const [deletedDocuments] = mockHardDeleteDocuments.mock.invocationCallOrder + expect(overdrawn).toBeLessThan(deletedDocuments) + expect(deletedDocuments).toBeLessThan(remaining) + expect(remaining).toBeLessThan(mockKnowledgeBaseContainerDelete.mock.invocationCallOrder[0]) + }) + it('soft-deletes abandoned KB bindings and removes their storage objects', async () => { dbChainMockFns.limit .mockResolvedValueOnce([{ key: 'kb/orphan-1' }, { key: 'kb/orphan-2' }]) diff --git a/apps/sim/background/cleanup-soft-deletes.ts b/apps/sim/background/cleanup-soft-deletes.ts index fa9ebe03925..ba955e7e2ab 100644 --- a/apps/sim/background/cleanup-soft-deletes.ts +++ b/apps/sim/background/cleanup-soft-deletes.ts @@ -40,6 +40,7 @@ import { resolveCleanupOwnerScope, } from '@/lib/cleanup/resource-scope' import { deduplicateFolderName } from '@/lib/folders/naming' +import { settleDetachedConnectorReservations } from '@/lib/knowledge/connectors/detachment' import { hardDeleteDocuments } from '@/lib/knowledge/documents/service' import type { StorageContext } from '@/lib/uploads' import { isUsingCloudStorage, StorageService } from '@/lib/uploads' @@ -445,11 +446,17 @@ async function cleanupExpiredKnowledgeBases( isNotNull(knowledgeBase.deletedAt), lt(knowledgeBase.deletedAt, retentionDate) ), - onBatch: (rows: { id: string }[]) => - hardDeleteKnowledgeBaseDocuments( - rows.map(({ id }) => id), - label - ), + /** + * The bases' DELETE cascades their connectors away, so a detached connector's reservation is + * settled here: an overdrawn one before the documents (their deletion floors usage at zero), + * the rest after them, so a deletion that fails partway leaves every step's ledger consistent. + */ + onBatch: async (rows: { id: string }[]) => { + const knowledgeBaseIds = rows.map(({ id }) => id) + await settleDetachedConnectorReservations(knowledgeBaseIds, 'overdrawn') + await hardDeleteKnowledgeBaseDocuments(knowledgeBaseIds, label) + await settleDetachedConnectorReservations(knowledgeBaseIds, 'remaining') + }, } return scope.kind === 'workspace' ? chunkedBatchDelete({ ...options, workspaceIds: scope.ids }) diff --git a/apps/sim/lib/knowledge/__integration__/drain-connector-event.ts b/apps/sim/lib/knowledge/__integration__/drain-connector-event.ts index 58e3ae5ccae..3697286f5fa 100644 --- a/apps/sim/lib/knowledge/__integration__/drain-connector-event.ts +++ b/apps/sim/lib/knowledge/__integration__/drain-connector-event.ts @@ -1,6 +1,6 @@ import { db } from '@sim/db' import { outboxEvent } from '@sim/db/schema' -import { and, eq, sql } from 'drizzle-orm' +import { and, asc, eq, sql } from 'drizzle-orm' import { expect } from 'vitest' import { processOutboxEventById } from '@/lib/core/outbox/service' import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler' @@ -13,9 +13,11 @@ export async function drainConnectorEvent(connectorId: string, eventType: string .where( and( eq(outboxEvent.eventType, eventType), + eq(outboxEvent.status, 'pending'), sql`${outboxEvent.payload}->>'connectorId' = ${connectorId}` ) ) + .orderBy(asc(outboxEvent.availableAt), asc(outboxEvent.id)) .limit(1) expect(job).toBeDefined() let status = await processOutboxEventById(job.id, knowledgeDocumentProcessingOutboxHandlers) diff --git a/apps/sim/lib/knowledge/__integration__/purged-detach-reservation.integration.ts b/apps/sim/lib/knowledge/__integration__/purged-detach-reservation.integration.ts new file mode 100644 index 00000000000..88d9c24e09d --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/purged-detach-reservation.integration.ts @@ -0,0 +1,355 @@ +/** + * Real PostgreSQL coverage for a knowledge base purged while one of its sources is still being + * detached. Removal charges the kept bytes up front as the source's reservation; the purge's + * cascade deletes the source, so the reservation has to be settled before it, or its bytes stay + * on the ledger for documents that no longer exist. + */ +import { db } from '@sim/db' +import { + document, + knowledgeBase, + knowledgeConnector, + organization, + outboxEvent, + user, + workspace, +} from '@sim/db/schema' +import { generateId } from '@sim/utils/id' +import { and, eq, inArray, isNull, sql } from 'drizzle-orm' +import { afterAll, describe, expect, it } from 'vitest' +import { processOutboxEventById } from '@/lib/core/outbox/service' +import { drainConnectorEvent } from '@/lib/knowledge/__integration__/drain-connector-event' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { + KNOWLEDGE_CONNECTOR_DETACH_EVENT, + settleDetachedConnectorReservations, +} from '@/lib/knowledge/connectors/detachment' +import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler' +import { createSingleDocument, hardDeleteDocuments } from '@/lib/knowledge/documents/service' +import { performDeleteKnowledgeConnector } from '@/lib/knowledge/orchestration/connectors' +import { runCleanupSoftDeletes } from '@/background/cleanup-soft-deletes' + +type Fixture = ReturnType +const fixtures: Fixture[] = [] + +/** More documents than one detach run releases (4 pages of 100), so a run can stop mid-source. */ +const SOURCE_DOCUMENTS = 450 +const RETENTION_HOURS = 720 + +async function seed() { + const ids = createKnowledgeAclFixtureIds() + fixtures.push(ids) + await seedKnowledgeAclFixture(ids) + await db + .update(knowledgeConnector) + .set({ accessMode: 'workspace' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + return ids +} + +function sourceDocument(ids: Fixture, bytes: number) { + return { + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + connectorId: ids.connectorId, + filename: 'source.txt', + fileUrl: 'data:text/plain;base64,c291cmNl', + fileSize: bytes, + mimeType: 'text/plain', + } +} + +async function manualDocument(ids: Fixture, bytes: number) { + return createSingleDocument( + { + filename: 'manual.txt', + fileUrl: `data:text/plain;base64,${Buffer.alloc(bytes, 'a').toString('base64')}`, + fileSize: bytes, + mimeType: 'text/plain', + }, + ids.knowledgeBaseId, + generateId(), + ids.aliceId + ) +} + +async function ledger(ids: Fixture) { + const [row] = await db + .select({ + workspaceBytes: workspace.storageUsedBytes, + payerBytes: organization.storageUsedBytes, + }) + .from(workspace) + .innerJoin(organization, eq(organization.id, workspace.organizationId)) + .where(eq(workspace.id, ids.workspaceId)) + return row +} + +async function reservation(ids: Fixture) { + const [row] = await db + .select({ reservedBytes: knowledgeConnector.detachReservedBytes }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + return row?.reservedBytes +} + +/** Removes the source keeping its documents; the release is left to the queued detach job. */ +async function disconnect(ids: Fixture) { + const outcome = await performDeleteKnowledgeConnector({ + knowledgeBase: { id: ids.knowledgeBaseId, name: 'Fixture', workspaceId: ids.workspaceId }, + connectorId: ids.connectorId, + deleteDocuments: false, + userId: ids.aliceId, + source: 'api', + requestId: generateId(), + recordSemanticAudit: false, + recordProductAnalytics: false, + }) + expect(outcome).toMatchObject({ success: true }) +} + +/** Runs the queued detach job exactly once, as one outbox worker pass would. */ +async function runDetachOnce(ids: Fixture) { + const [job] = await db + .select({ id: outboxEvent.id }) + .from(outboxEvent) + .where( + and( + eq(outboxEvent.eventType, KNOWLEDGE_CONNECTOR_DETACH_EVENT), + eq(outboxEvent.status, 'pending'), + sql`${outboxEvent.payload}->>'connectorId' = ${ids.connectorId}` + ) + ) + .limit(1) + return processOutboxEventById(job.id, knowledgeDocumentProcessingOutboxHandlers) +} + +afterAll(async () => { + for (const ids of fixtures) { + await db + .delete(outboxEvent) + .where( + and( + eq(outboxEvent.eventType, KNOWLEDGE_CONNECTOR_DETACH_EVENT), + sql`${outboxEvent.payload}->>'knowledgeBaseId' = ${ids.knowledgeBaseId}` + ) + ) + await db.delete(knowledgeBase).where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db.delete(workspace).where(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])) + } + await db.$client.end() +}) + +describe('purging a knowledge base with a source still being detached', () => { + it('settles the reservation before the purge and leaves the late detach job nothing to settle', async () => { + const ids = await seed() + const rows = Array.from({ length: SOURCE_DOCUMENTS }, (_, index) => + sourceDocument(ids, (index % 7) + 1) + ) + const keptBytes = rows.reduce((total, row) => total + row.fileSize, 0) + await db.insert(document).values(rows) + + await disconnect(ids) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + expect(await reservation(ids)).toBe(keptBytes) + + /** One bounded run releases some documents, which consume their share of the reservation. */ + expect(await runDetachOnce(ids)).toBe('pending') + const [released] = await db + .select({ bytes: sql`COALESCE(SUM(${document.fileSize}), 0)::integer` }) + .from(document) + .where(and(eq(document.knowledgeBaseId, ids.knowledgeBaseId), isNull(document.connectorId))) + expect(released.bytes).toBeGreaterThan(0) + expect(released.bytes).toBeLessThan(keptBytes) + expect(await reservation(ids)).toBe(keptBytes - released.bytes) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + await db + .update(knowledgeBase) + .set({ deletedAt: new Date(Date.now() - 2 * RETENTION_HOURS * 60 * 60 * 1000) }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await runCleanupSoftDeletes({ + label: 'purge-fixture', + plan: 'free', + retentionHours: RETENTION_HOURS, + workspaceIds: [ids.workspaceId], + }) + + expect( + await db + .select({ id: knowledgeBase.id }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + ).toHaveLength(0) + expect( + await db + .select({ id: knowledgeConnector.id }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + ).toHaveLength(0) + /** The workspace keeps no files or documents, so a from-scratch recount is zero. */ + expect(await ledger(ids)).toEqual({ workspaceBytes: 0, payerBytes: 0 }) + + await drainConnectorEvent(ids.connectorId, KNOWLEDGE_CONNECTOR_DETACH_EVENT) + expect(await ledger(ids)).toEqual({ workspaceBytes: 0, payerBytes: 0 }) + }) + + it('settles an overdrawn reservation before the purge deletes the released documents', async () => { + const ids = await seed() + const rows = Array.from({ length: SOURCE_DOCUMENTS }, (_, index) => + sourceDocument(ids, (index % 7) + 1) + ) + const keptBytes = rows.reduce((total, row) => total + row.fileSize, 0) + await db.insert(document).values(rows) + await disconnect(ids) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + /** Documents that grew after removal release more bytes than removal reserved. */ + await db + .update(document) + .set({ fileSize: sql`${document.fileSize} + 10` }) + .where(eq(document.connectorId, ids.connectorId)) + expect(await runDetachOnce(ids)).toBe('pending') + const overdrawn = await reservation(ids) + expect(overdrawn).toBeLessThan(0) + /** Removal charged less than the released documents now hold, so the ledger trails them. */ + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + await db + .update(knowledgeBase) + .set({ deletedAt: new Date(Date.now() - 2 * RETENTION_HOURS * 60 * 60 * 1000) }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await runCleanupSoftDeletes({ + label: 'purge-overdrawn-fixture', + plan: 'free', + retentionHours: RETENTION_HOURS, + workspaceIds: [ids.workspaceId], + }) + + expect(await ledger(ids)).toEqual({ workspaceBytes: 0, payerBytes: 0 }) + }) + + it('keeps a positive reservation billed and the detach running when the purge stops before the documents', async () => { + const ids = await seed() + const rows = Array.from({ length: SOURCE_DOCUMENTS }, (_, index) => + sourceDocument(ids, (index % 7) + 1) + ) + const keptBytes = rows.reduce((total, row) => total + row.fileSize, 0) + await db.insert(document).values(rows) + await disconnect(ids) + + /** The purge settled only overdrawn reservations, then its document deletion failed. */ + await settleDetachedConnectorReservations([ids.knowledgeBaseId], 'overdrawn') + expect(await reservation(ids)).toBe(keptBytes) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + /** The detach is not fenced off, so a base restored now still releases its documents. */ + expect(await runDetachOnce(ids)).toBe('pending') + const [released] = await db + .select({ bytes: sql`COALESCE(SUM(${document.fileSize}), 0)::integer` }) + .from(document) + .where(and(eq(document.knowledgeBaseId, ids.knowledgeBaseId), isNull(document.connectorId))) + expect(released.bytes).toBeGreaterThan(0) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + /** A retried purge finishes the job and leaves the ledger at a from-scratch recount. */ + await db + .update(knowledgeBase) + .set({ deletedAt: new Date(Date.now() - 2 * RETENTION_HOURS * 60 * 60 * 1000) }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await runCleanupSoftDeletes({ + label: 'purge-retry-fixture', + plan: 'free', + retentionHours: RETENTION_HOURS, + workspaceIds: [ids.workspaceId], + }) + expect(await ledger(ids)).toEqual({ workspaceBytes: 0, payerBytes: 0 }) + }) + + it('pauses the release while the base is deleted and resumes it after a restore', async () => { + const ids = await seed() + const rows = Array.from({ length: SOURCE_DOCUMENTS }, (_, index) => + sourceDocument(ids, (index % 7) + 1) + ) + const keptBytes = rows.reduce((total, row) => total + row.fileSize, 0) + await db.insert(document).values(rows) + await disconnect(ids) + const releasedDocuments = async () => { + const [row] = await db + .select({ count: sql`count(*)::integer` }) + .from(document) + .where(and(eq(document.knowledgeBaseId, ids.knowledgeBaseId), isNull(document.connectorId))) + return row.count + } + + await db + .update(knowledgeBase) + .set({ deletedAt: new Date() }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + expect(await runDetachOnce(ids)).toBe('pending') + expect(await releasedDocuments()).toBe(0) + expect(await reservation(ids)).toBe(keptBytes) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + + await db + .update(knowledgeBase) + .set({ deletedAt: null }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + /** The paused event rechecks an hour later; the worker would pick it up then. */ + await db + .update(outboxEvent) + .set({ availableAt: new Date() }) + .where( + and( + eq(outboxEvent.eventType, KNOWLEDGE_CONNECTOR_DETACH_EVENT), + sql`${outboxEvent.payload}->>'connectorId' = ${ids.connectorId}` + ) + ) + expect(await runDetachOnce(ids)).toBe('pending') + expect(await releasedDocuments()).toBeGreaterThan(0) + expect(await ledger(ids)).toEqual({ workspaceBytes: keptBytes, payerBytes: keptBytes }) + }) + + it('zeroes a refunded reservation so a detach run that still reaches the source refunds nothing', async () => { + const ids = await seed() + await manualDocument(ids, 29) + const source = sourceDocument(ids, 37) + await db.insert(document).values(source) + + await disconnect(ids) + expect(await ledger(ids)).toEqual({ workspaceBytes: 66, payerBytes: 66 }) + expect(await hardDeleteDocuments([source.id], generateId())).toBe(1) + + await settleDetachedConnectorReservations([ids.knowledgeBaseId], 'remaining') + expect(await reservation(ids)).toBe(0) + expect(await ledger(ids)).toEqual({ workspaceBytes: 29, payerBytes: 29 }) + await settleDetachedConnectorReservations([ids.knowledgeBaseId], 'remaining') + expect(await ledger(ids)).toEqual({ workspaceBytes: 29, payerBytes: 29 }) + + await drainConnectorEvent(ids.connectorId, KNOWLEDGE_CONNECTOR_DETACH_EVENT) + expect(await ledger(ids)).toEqual({ workspaceBytes: 29, payerBytes: 29 }) + }) + + it('charges an overdrawn reservation once, as the detach job would', async () => { + const ids = await seed() + await manualDocument(ids, 29) + await disconnect(ids) + /** Released bytes beyond what removal charged, e.g. a document that grew before its release. */ + await db + .update(knowledgeConnector) + .set({ detachReservedBytes: -7 }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + + await settleDetachedConnectorReservations([ids.knowledgeBaseId], 'remaining') + expect(await reservation(ids)).toBe(0) + expect(await ledger(ids)).toEqual({ workspaceBytes: 36, payerBytes: 36 }) + + await drainConnectorEvent(ids.connectorId, KNOWLEDGE_CONNECTOR_DETACH_EVENT) + expect(await ledger(ids)).toEqual({ workspaceBytes: 36, payerBytes: 36 }) + }) +}) diff --git a/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts b/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts index c5474d349a2..163f1fcea0d 100644 --- a/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts @@ -27,8 +27,9 @@ import { isRecordLike } from '@sim/utils/object' import { eq, inArray, sql } from 'drizzle-orm' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +/** The TINQL `resolveTinKeywordQuery` renders for `fixture`: its `english` stem, quoted. */ vi.mock('@/lib/knowledge/search/tin-keyword', () => ({ - resolveTinKeywordQuery: async () => 'fixture', + resolveTinKeywordQuery: async () => '"fixtur"', })) import { @@ -284,13 +285,6 @@ beforeAll(async () => { embeddingModel: 'text-embedding-3-small', embedding: [1, ...Array(1535).fill(0)], }) - await db.insert(embeddingKeywordTin).values({ - id: embeddingId, - knowledgeBaseId: ids.knowledgeBaseId, - documentId, - enabled: true, - content: 'fixture readme', - }) const [tin] = await db.execute<{ present: boolean }>( sql`SELECT to_regnamespace('tin') IS NOT NULL AS present` ) @@ -300,10 +294,21 @@ beforeAll(async () => { sql.raw(`CREATE SCHEMA tin; CREATE FUNCTION tin.full_score(tid) RETURNS double precision LANGUAGE sql IMMUTABLE AS 'SELECT 1.0::float8'; CREATE FUNCTION knowledge_tin_base_token(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$SELECT 'kb'$$; + CREATE FUNCTION knowledge_tin_stream(vector tsvector) RETURNS text LANGUAGE sql IMMUTABLE AS $$ + SELECT coalesce(string_agg(entry.lexeme, ' ' ORDER BY position), '') + FROM unnest(vector) AS entry(lexeme, positions, weights), unnest(entry.positions) AS position + $$; CREATE FUNCTION tin_fixture_match(text, text) RETURNS boolean LANGUAGE sql IMMUTABLE AS 'SELECT true'; CREATE OPERATOR ==> (LEFTARG = text, RIGHTARG = text, FUNCTION = tin_fixture_match);`) ) } + /** Written as the projection trigger writes it, so real Tin scopes the row to its base. */ + await db.execute(sql` + INSERT INTO ${embeddingKeywordTin} (id, knowledge_base_id, document_id, enabled, content) + SELECT id, knowledge_base_id, document_id, enabled, + knowledge_tin_base_token(knowledge_base_id) || ' ' || knowledge_tin_stream(content_tsv) + FROM ${embedding} WHERE id = ${embeddingId} + ON CONFLICT (id) DO UPDATE SET content = EXCLUDED.content`) }) afterAll(async () => { @@ -312,6 +317,7 @@ afterAll(async () => { sql.raw(`DROP OPERATOR IF EXISTS ==> (text, text); DROP FUNCTION IF EXISTS tin_fixture_match(text, text); DROP FUNCTION IF EXISTS knowledge_tin_base_token(text); + DROP FUNCTION IF EXISTS knowledge_tin_stream(tsvector); DROP SCHEMA IF EXISTS tin CASCADE;`) ) } diff --git a/apps/sim/lib/knowledge/connectors/detachment.test.ts b/apps/sim/lib/knowledge/connectors/detachment.test.ts index 13ce1a15acb..6d04b625d00 100644 --- a/apps/sim/lib/knowledge/connectors/detachment.test.ts +++ b/apps/sim/lib/knowledge/connectors/detachment.test.ts @@ -36,6 +36,7 @@ import { detachKnowledgeConnector, enqueueConnectorDetachment, KNOWLEDGE_CONNECTOR_DETACH_EVENT, + settleDetachedConnectorReservations, } from '@/lib/knowledge/connectors/detachment' const payload = { @@ -210,6 +211,17 @@ describe('connector detachment', () => { expect(mocks.revoke).not.toHaveBeenCalled() }) + it('releases nothing and spends no attempt while the knowledge base is deleted', async () => { + queueTableRows(knowledgeBase, [{ ...owner, deletedAt: new Date('2026-09-20T00:00:00.000Z') }]) + + const result = await detachKnowledgeConnector(payload, context()) + + expect(result).toMatchObject({ outcome: 'deferred', consumeAttempt: false }) + expect(dbChainMockFns.update).not.toHaveBeenCalled() + expect(mocks.decrementStorage).not.toHaveBeenCalled() + expect(mocks.revoke).not.toHaveBeenCalled() + }) + it('stops when the knowledge base is gone', async () => { resetDbChainMock() await detachKnowledgeConnector(payload, context()) @@ -225,3 +237,46 @@ describe('connector detachment', () => { expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) }) + +describe('purged knowledge base reservations', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + mocks.resolveStorage.mockResolvedValue(STORAGE_CONTEXT) + mocks.incrementStorage.mockResolvedValue(1_000) + /** The base is read once to resolve its payer and again under the settlement lock. */ + queueTableRows(knowledgeBase, [owner]) + queueTableRows(knowledgeBase, [owner]) + }) + afterEach(resetDbChainMock) + + it('settles a base as one net refund, so no overdraft warning precedes it', async () => { + queueTableRows(knowledgeConnector, [ + { id: 'connector-a', reservedBytes: -7 }, + { id: 'connector-b', reservedBytes: 50 }, + ]) + + await settleDetachedConnectorReservations(['kb-1'], 'remaining') + + expect(mocks.decrementStorage).toHaveBeenCalledOnce() + expect(mocks.decrementStorage).toHaveBeenCalledWith(expect.anything(), STORAGE_CONTEXT, 43) + expect(mocks.incrementStorage).not.toHaveBeenCalled() + expect(mocks.notifyStorage).not.toHaveBeenCalled() + /** Settlement zeroes the reservation and leaves the detach itself untouched. */ + expect(dbChainMockFns.set).toHaveBeenCalledWith({ detachReservedBytes: 0 }) + }) + + it('charges a net overdraft once and notifies with the final balance', async () => { + queueTableRows(knowledgeConnector, [ + { id: 'connector-a', reservedBytes: -30 }, + { id: 'connector-b', reservedBytes: 10 }, + ]) + + await settleDetachedConnectorReservations(['kb-1'], 'remaining') + + expect(mocks.incrementStorage).toHaveBeenCalledOnce() + expect(mocks.incrementStorage).toHaveBeenCalledWith(expect.anything(), STORAGE_CONTEXT, 20) + expect(mocks.decrementStorage).not.toHaveBeenCalled() + expect(mocks.notifyStorage).toHaveBeenCalledWith(STORAGE_CONTEXT, 1_000) + }) +}) diff --git a/apps/sim/lib/knowledge/connectors/detachment.ts b/apps/sim/lib/knowledge/connectors/detachment.ts index e93feb1f006..fee3c8e4d04 100644 --- a/apps/sim/lib/knowledge/connectors/detachment.ts +++ b/apps/sim/lib/knowledge/connectors/detachment.ts @@ -6,7 +6,7 @@ import { knowledgeBase, knowledgeConnector, } from '@sim/db/schema' -import { and, eq, inArray, isNotNull, sql } from 'drizzle-orm' +import { and, asc, eq, inArray, isNotNull, lt, ne, sql } from 'drizzle-orm' import { z } from 'zod' import { decrementStorageUsageForBillingContextInTx, @@ -17,6 +17,7 @@ import { } from '@/lib/billing/storage' import { continueOutboxHandler, + deferOutboxHandler, enqueueOutboxEvent, type OutboxHandler, } from '@/lib/core/outbox/service' @@ -29,6 +30,8 @@ const DOCUMENT_BATCH_SIZE = 100 const PROJECTION_ROW_BATCH_SIZE = 250 const MAX_BATCHES_PER_RUN = 4 const RUN_BUDGET_MS = 30_000 +/** How often a detachment paused on a deleted knowledge base checks for its restore or purge. */ +const DELETED_BASE_RECHECK_MS = 60 * 60 * 1000 const detachmentPayloadSchema = z .object({ @@ -77,6 +80,118 @@ export function keptDocumentBytes() { const SEARCH_PROJECTIONS = [embeddingSearch, embeddingKeywordTin] as const +/** + * Settles what is left of a detached connector's reservation: an unreleased remainder is refunded, + * and an overdraft, released bytes beyond what removal charged, is charged as already admitted. + * Returns the payer's updated usage when it grew, for a storage-limit notification after commit. + */ +async function settleDetachReservationInTx( + tx: DbOrTx, + storageContext: StorageBillingContext, + reservedBytes: number +): Promise { + if (reservedBytes > 0) { + await decrementStorageUsageForBillingContextInTx(tx, storageContext, reservedBytes) + return undefined + } + if (reservedBytes < 0) { + return incrementAdmittedStorageUsageForBillingContextInTx(tx, storageContext, -reservedBytes) + } + return undefined +} + +/** + * Which detach reservations a purge settles at each of its two points. + * `overdrawn` settles only negative reservations and runs before the documents are deleted; + * `remaining` settles whatever is left and runs after them. + */ +export type DetachReservationSettlement = 'overdrawn' | 'remaining' + +/** + * Settles the reservations of detached connectors on knowledge bases being hard-deleted. + * + * Purging a base cascades its connectors away, and with them the reservation a detached connector + * still holds for documents it never released; its pending detach job then finds no base and + * settles nothing. So the purge settles them itself, locking each base's detached connectors in + * the detach job's order (base, then connector), settling their net reservation exactly as the + * job's final transaction would, and zeroing it in the same transaction so nothing settles twice. + * + * The ledger must match the base's documents after every step, since document deletion can fail + * partway and be retried. Deleting a released document decrements usage with a floor at zero, so + * an overdrawn reservation settled afterwards would re-add bytes the floor discarded: `overdrawn` + * settles those before the documents go, which only charges bytes the released documents already + * hold. A positive reservation still pays for documents that remain until they are deleted, so + * `remaining` settles it afterwards. No detach page interleaves: a detach job releases nothing while its + * base is deleted, and a base restored before the purge completes resumes its detach unchanged. + */ +export async function settleDetachedConnectorReservations( + knowledgeBaseIds: string[], + settlement: DetachReservationSettlement +): Promise { + for (const knowledgeBaseId of knowledgeBaseIds) { + const [owner] = await db + .select({ workspaceId: knowledgeBase.workspaceId }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, knowledgeBaseId)) + .limit(1) + if (!owner?.workspaceId) continue + const storageContext = await resolveStorageBillingContext(owner.workspaceId) + + const updatedUsage = await db.transaction(async (tx) => { + await tx.execute(sql`SET LOCAL lock_timeout = '5s'`) + await tx.execute(sql`SET LOCAL statement_timeout = '30s'`) + const [lockedOwner] = await tx + .select({ workspaceId: knowledgeBase.workspaceId }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, knowledgeBaseId)) + .for('share') + .limit(1) + if (!lockedOwner) return undefined + if (lockedOwner.workspaceId !== owner.workspaceId) { + throw new Error('Knowledge base workspace changed during detach reservation settlement') + } + const reserved = await tx + .select({ + id: knowledgeConnector.id, + reservedBytes: knowledgeConnector.detachReservedBytes, + }) + .from(knowledgeConnector) + .where( + and( + eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId), + isNotNull(knowledgeConnector.detachedAt), + settlement === 'overdrawn' + ? lt(knowledgeConnector.detachReservedBytes, 0) + : ne(knowledgeConnector.detachReservedBytes, 0) + ) + ) + .orderBy(asc(knowledgeConnector.id)) + .for('update') + if (reserved.length === 0) return undefined + + /** + * One net settlement per base: the ledger lands where settling each connector in turn would + * leave it, and the notifier sees that final balance rather than one from mid-sequence. + */ + const netReservedBytes = reserved.reduce((sum, connector) => sum + connector.reservedBytes, 0) + const grownUsage = await settleDetachReservationInTx(tx, storageContext, netReservedBytes) + await tx + .update(knowledgeConnector) + .set({ detachReservedBytes: 0 }) + .where( + inArray( + knowledgeConnector.id, + reserved.map(({ id }) => id) + ) + ) + return grownUsage + }) + if (updatedUsage !== undefined) { + await maybeNotifyStorageLimitForBillingContext(storageContext, updatedUsage) + } + } +} + /** * Releases a detached connector's documents as standalone entries, then deletes the connector. * @@ -110,7 +225,7 @@ export const detachKnowledgeConnector: OutboxHandler = async (rawPayload, contex : undefined let storageNotification: { context: StorageBillingContext; updatedUsage: number } | undefined - let outcome: 'progress' | 'complete' | 'obsolete' = 'progress' + let outcome: 'progress' | 'complete' | 'obsolete' | 'paused' = 'progress' for (let batch = 0; batch < MAX_BATCHES_PER_RUN && outcome === 'progress'; batch++) { context.signal.throwIfAborted() outcome = await db.transaction(async (tx) => { @@ -118,7 +233,7 @@ export const detachKnowledgeConnector: OutboxHandler = async (rawPayload, contex await tx.execute(sql`SET LOCAL statement_timeout = '30s'`) /** Match source writes and document deletion: parent KB, connector, then documents. */ const [lockedOwner] = await tx - .select({ workspaceId: knowledgeBase.workspaceId }) + .select({ workspaceId: knowledgeBase.workspaceId, deletedAt: knowledgeBase.deletedAt }) .from(knowledgeBase) .where(eq(knowledgeBase.id, payload.knowledgeBaseId)) .for('share') @@ -127,6 +242,13 @@ export const detachKnowledgeConnector: OutboxHandler = async (rawPayload, contex if (lockedOwner.workspaceId !== owner.workspaceId) { throw new Error('Knowledge base workspace changed during connector detachment') } + /** + * A deleted base releases nothing: its documents stay archived and attached, and the + * reservation keeps paying for them, until a restore resumes the release or the purge + * settles it. Checked under the base's share lock, so no page releases once the deletion + * commits, and none can interleave with the purge's settlement and document deletion. + */ + if (lockedOwner.deletedAt) return 'paused' const [connector] = await tx .select({ detachedAt: knowledgeConnector.detachedAt, @@ -165,21 +287,13 @@ export const detachKnowledgeConnector: OutboxHandler = async (rawPayload, contex context.signal ) if (drained === 'complete' && storageContext) { - if (connector.reservedBytes > 0) { - await decrementStorageUsageForBillingContextInTx( - tx, - storageContext, - connector.reservedBytes - ) - } else if (connector.reservedBytes < 0) { - const updatedUsage = await incrementAdmittedStorageUsageForBillingContextInTx( - tx, - storageContext, - -connector.reservedBytes - ) - if (updatedUsage !== undefined) { - storageNotification = { context: storageContext, updatedUsage } - } + const updatedUsage = await settleDetachReservationInTx( + tx, + storageContext, + connector.reservedBytes + ) + if (updatedUsage !== undefined) { + storageNotification = { context: storageContext, updatedUsage } } } return drained @@ -241,6 +355,17 @@ export const detachKnowledgeConnector: OutboxHandler = async (rawPayload, contex ) } if (outcome === 'obsolete') return + if (outcome === 'paused') { + /** + * Waiting spends no attempt: the event still reaches a terminal state, since either a + * restore resumes the release or the purge removes the base and the next run completes. + */ + return deferOutboxHandler( + 'Connector detachment waits while its knowledge base is deleted', + DELETED_BASE_RECHECK_MS, + false + ) + } if (outcome === 'complete') { if (payload.credentialAccess) { context.signal.throwIfAborted() diff --git a/apps/sim/lib/memory/checkpoint-codec.test.ts b/apps/sim/lib/memory/checkpoint-codec.test.ts index 4d71765faa8..31ab069b8cf 100644 --- a/apps/sim/lib/memory/checkpoint-codec.test.ts +++ b/apps/sim/lib/memory/checkpoint-codec.test.ts @@ -73,7 +73,10 @@ describe('private memory checkpoint encoding', () => { it('rejects altered ciphertext and oversized plaintext before decryption', async () => { const encrypted = await encryptMemoryCheckpoint({ private: 'value' }) - await expect(decryptMemoryCheckpoint(`${encrypted.slice(0, -2)}ff`)).rejects.toThrow() + /** Always change the auth tag's last hex digit; overwriting it with a fixed value can be a no-op. */ + const tampered = `${encrypted.slice(0, -1)}${encrypted.endsWith('0') ? '1' : '0'}` + expect(tampered).not.toBe(encrypted) + await expect(decryptMemoryCheckpoint(tampered)).rejects.toThrow() await expect(encryptMemoryCheckpoint({ output: 'a'.repeat(3 * 1024 * 1024) })).rejects.toThrow( 'byte limit' ) diff --git a/scripts/test-knowledge-acls.ts b/scripts/test-knowledge-acls.ts index b7f0d4bc030..6b43057b550 100644 --- a/scripts/test-knowledge-acls.ts +++ b/scripts/test-knowledge-acls.ts @@ -203,7 +203,12 @@ try { ) run( 'bunx', - ['vitest', 'run', 'script-migrations/0016_backfill_search_vectors.postgres.test.ts'], + [ + 'vitest', + 'run', + 'script-migrations/0016_backfill_search_vectors.postgres.test.ts', + 'script-migrations/0021_embedding_search_connector.postgres.test.ts', + ], { cwd: path.join(root, 'packages/db'), env: environment,