fix(knowledge): carry detach reservations through payer moves and refuse credentials to removed connectors

Author: waleedlatif1

apps/sim/lib/billing/storage/payer-transfer.test.ts

@@ -31,6 +31,12 @@ vi.mock('@sim/db/schema', () => ({
     id: 'knowledgeBase.id',
     workspaceId: 'knowledgeBase.workspaceId',
   },
+  knowledgeConnector: {
+    __table: 'knowledgeConnector',
+    detachedAt: 'knowledgeConnector.detachedAt',
+    detachReservedBytes: 'knowledgeConnector.detachReservedBytes',
+    knowledgeBaseId: 'knowledgeConnector.knowledgeBaseId',
+  },
   organization: {
     __table: 'organization',
     id: 'organization.id',
@@ -535,7 +541,10 @@ describe('changeWorkspaceStoragePayerInTx', () => {
     expect(query.values).not.toContain('workspaceFiles.deletedAt')
     expect(query.values).toContain('document.connectorId')
     expect(query.values).toContain('document.deletedAt')
-    expect(query.values.filter((value) => value === 'workspace-1')).toHaveLength(3)
+    /** A detaching connector's reservation is already charged, so a payer move carries it. */
+    expect(query.values).toContain('knowledgeConnector.detachReservedBytes')
+    expect(query.values).toContain('knowledgeConnector.detachedAt')
+    expect(query.values.filter((value) => value === 'workspace-1')).toHaveLength(4)
   })
 
   it('fails closed when a billable file is missing canonical size metadata', async () => {

apps/sim/lib/billing/storage/payer-transfer.ts

@@ -1,6 +1,7 @@
 import {
   document,
   knowledgeBase,
+  knowledgeConnector,
   organization,
   userStats,
   workspace,
@@ -77,7 +78,8 @@ function parseExactBytes(value: number | string, label: string): number {
  * Computes one workspace's live billable bytes with two index-bounded scalar
  * aggregates. Archived workspace files and documents remain billable while
  * their objects are retained; mothership files, connector documents, and
- * deleted documents are excluded.
+ * deleted documents are excluded. A detaching connector's reservation counts:
+ * it was charged when the connector was removed with its documents kept.
  */
 async function getExactWorkspaceStorageBytes(tx: DbOrTx, workspaceId: string): Promise<number> {
   const [row] = await tx.execute<ExactWorkspaceStorageRow>(sql`
@@ -103,6 +105,13 @@ async function getExactWorkspaceStorageBytes(tx: DbOrTx, workspaceId: string): P
         WHERE ${knowledgeBase.workspaceId} = ${workspaceId}
           AND ${document.connectorId} IS NULL
           AND ${document.deletedAt} IS NULL
+      ), 0)::bigint + COALESCE((
+        SELECT SUM(${knowledgeConnector.detachReservedBytes})
+        FROM ${knowledgeConnector}
+        INNER JOIN ${knowledgeBase}
+          ON ${knowledgeBase.id} = ${knowledgeConnector.knowledgeBaseId}
+        WHERE ${knowledgeBase.workspaceId} = ${workspaceId}
+          AND ${knowledgeConnector.detachedAt} IS NOT NULL
       ), 0)::bigint AS document_bytes
   `)
 
@@ -209,6 +218,20 @@ async function getExactWorkspaceStorageBytesBatch(
         AND ${document.connectorId} IS NULL
         AND ${document.deletedAt} IS NULL
       GROUP BY ${knowledgeBase.workspaceId}
+
+      UNION ALL
+
+      SELECT
+        ${knowledgeBase.workspaceId} AS workspace_id,
+        0::bigint AS workspace_file_bytes,
+        SUM(${knowledgeConnector.detachReservedBytes}) AS document_bytes,
+        0::bigint AS workspace_file_missing_size_count
+      FROM ${knowledgeConnector}
+      INNER JOIN ${knowledgeBase}
+        ON ${knowledgeBase.id} = ${knowledgeConnector.knowledgeBaseId}
+      WHERE ${inArray(knowledgeBase.workspaceId, workspaceIds)}
+        AND ${knowledgeConnector.detachedAt} IS NOT NULL
+      GROUP BY ${knowledgeBase.workspaceId}
     ) storage_by_workspace
     GROUP BY storage_by_workspace.workspace_id
     ORDER BY storage_by_workspace.workspace_id

apps/sim/lib/knowledge/connectors/member-access.test.ts

@@ -2,7 +2,7 @@
  * @vitest-environment node
  */
 import { credentialGroup, knowledgeBase, knowledgeConnector } from '@sim/db/schema'
-import { dbChainMockFns, hasMockCondition, resetDbChainMock } from '@sim/testing'
+import { dbChainMockFns, hasMockCondition, queueTableRows, resetDbChainMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
@@ -307,6 +307,7 @@ describe('knowledge connector member access', () => {
     })
 
     it('resolves the token when the policy names the connector under the credential option and audits it', async () => {
+      queueTableRows(knowledgeConnector, [{ id: 'connector-1' }])
       mocks.requireResourcePolicy.mockResolvedValue(
         storedPolicy(2, [
           { credentialGroupOptionId: 'option-drive', connectorIds: ['connector-1'] },
@@ -338,7 +339,21 @@ describe('knowledge connector member access', () => {
       )
     })
 
+    it('refuses a removed connector even while its policy grant remains', async () => {
+      mocks.requireResourcePolicy.mockResolvedValue(
+        storedPolicy(2, [
+          { credentialGroupOptionId: 'option-drive', connectorIds: ['connector-1'] },
+        ])
+      )
+
+      await expect(mintKnowledgeConnectorMemberToken(mintInput)).rejects.toThrow(
+        'Knowledge connector has been removed'
+      )
+      expect(mocks.resolveManagedOAuthToken).not.toHaveBeenCalled()
+    })
+
     it('reports provider rejection only under the current connector grant', async () => {
+      queueTableRows(knowledgeConnector, [{ id: 'connector-1' }])
       mocks.requireResourcePolicy.mockResolvedValue(
         storedPolicy(2, [
           { credentialGroupOptionId: 'option-drive', connectorIds: ['connector-1'] },
@@ -451,6 +466,7 @@ describe('knowledge connector member access', () => {
 
   describe('list', () => {
     it('pages the option credentials only for a granted connector', async () => {
+      queueTableRows(knowledgeConnector, [{ id: 'connector-1' }])
       mocks.requireResourcePolicy.mockResolvedValue(
         storedPolicy(2, [
           { credentialGroupOptionId: 'option-drive', connectorIds: ['connector-1'] },

apps/sim/lib/knowledge/connectors/member-access.ts

@@ -302,6 +302,15 @@ export async function assertKnowledgeConnectorCredentialAccess(
     }
     return
   }
+  /** The policy grant outlives a removal until its revocation lands; the connector row does not. */
+  const [liveConnector] = await db
+    .select({ id: knowledgeConnector.id })
+    .from(knowledgeConnector)
+    .where(and(eq(knowledgeConnector.id, binding.connectorId), connectorIsLive()))
+    .limit(1)
+  if (!liveConnector) {
+    throw new KnowledgeConnectorMemberAccessDeniedError('Knowledge connector has been removed')
+  }
   const policy = await requireResourcePolicy(
     policyTarget({ ...binding, workspaceId: scope.workspaceId })
   ).catch((error: unknown) => {