fix billing account details for kb embeddings

Author: icecrasher321

apps/sim/app/api/auth/sso/providers/route.ts

@@ -13,7 +13,7 @@ const logger = createLogger('SSOProvidersRoute')
 export const GET = withRouteHandler(async (request: NextRequest) => {
   try {
     const session = await getSession()
-    const parsed = await parseRequest(listSsoProvidersContract, request, undefined)
+    const parsed = await parseRequest(listSsoProvidersContract, request, {})
     if (!parsed.success) return parsed.response
     const { organizationId } = parsed.data.query
 

apps/sim/app/api/wand/route.ts

@@ -14,6 +14,7 @@ import { getCostMultiplier, isBillingEnabled } from '@/lib/core/config/feature-f
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { enrichTableSchema } from '@/lib/table/llm/wand'
+import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
 import { verifyWorkspaceMembership } from '@/app/api/workflows/utils'
 import { extractResponseText, parseResponsesUsage } from '@/providers/openai/utils'
 import { getModelPricing } from '@/providers/utils'
@@ -86,7 +87,8 @@ Use this context to calculate relative dates like "yesterday", "last week", "beg
 }
 
 async function updateUserStatsForWand(
-  userId: string,
+  billingUserId: string,
+  workspaceId: string | null,
   usage: {
     prompt_tokens?: number
     completion_tokens?: number
@@ -128,7 +130,8 @@ async function updateUserStatsForWand(
     }
 
     await recordUsage({
-      userId,
+      userId: billingUserId,
+      workspaceId: workspaceId ?? undefined,
       entries: [
         {
           category: 'model',
@@ -143,7 +146,7 @@ async function updateUserStatsForWand(
       },
     })
 
-    await checkAndBillOverageThreshold(userId)
+    await checkAndBillOverageThreshold(billingUserId)
   } catch (error) {
     logger.error(`[${requestId}] Failed to update user stats for wand usage`, error)
   }
@@ -223,6 +226,21 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
       }
     }
 
+    let billingUserId = session.user.id
+    if (workspaceId) {
+      const workspaceBilledAccountUserId = await getWorkspaceBilledAccountUserId(workspaceId)
+      if (!workspaceBilledAccountUserId) {
+        logger.error(`[${requestId}] Unable to resolve billed account for workspace`, {
+          workspaceId,
+        })
+        return NextResponse.json(
+          { success: false, error: 'Unable to resolve billing account for this workspace' },
+          { status: 500 }
+        )
+      }
+      billingUserId = workspaceBilledAccountUserId
+    }
+
     let isBYOK = false
     let activeOpenAIKey = openaiApiKey
 
@@ -339,7 +357,13 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
               }
 
               usageRecorded = true
-              await updateUserStatsForWand(session.user.id, finalUsage, requestId, isBYOK)
+              await updateUserStatsForWand(
+                billingUserId,
+                workspaceId,
+                finalUsage,
+                requestId,
+                isBYOK
+              )
             }
 
             try {
@@ -556,7 +580,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
     const usage = parseResponsesUsage(completion.usage)
     if (usage) {
       await updateUserStatsForWand(
-        session.user.id,
+        billingUserId,
+        workspaceId,
         {
           prompt_tokens: usage.promptTokens,
           completion_tokens: usage.completionTokens,

apps/sim/app/workspace/[workspaceId]/w/hooks/use-import-workspace.ts

@@ -3,11 +3,11 @@ import { createLogger } from '@sim/logger'
 import { generateId } from '@sim/utils/id'
 import { useRouter } from 'next/navigation'
 import { requestJson } from '@/lib/api/client/request'
-import type { ContractBodyInput } from '@/lib/api/contracts'
 import {
   createWorkflowContract,
   createWorkspaceContract,
   putWorkflowNormalizedStateContract,
+  type WorkflowStateContractInput,
   workflowVariablesContract,
 } from '@/lib/api/contracts'
 import {
@@ -186,10 +186,31 @@ export function useImportWorkspace({ onSuccess }: UseImportWorkspaceProps = {})
               continue
             }
 
+            type ContractEdgeInput = WorkflowStateContractInput['edges'][number]
+
+            const sanitizedEdges: ContractEdgeInput[] = (workflowData.edges || []).map((edge) => {
+              const { sourceHandle, targetHandle, ...rest } = edge
+              const sanitized: ContractEdgeInput = { ...rest } as ContractEdgeInput
+              if (typeof sourceHandle === 'string' && sourceHandle.length > 0) {
+                sanitized.sourceHandle = sourceHandle
+              }
+              if (typeof targetHandle === 'string' && targetHandle.length > 0) {
+                sanitized.targetHandle = targetHandle
+              }
+              return sanitized
+            })
+
+            const workflowStateBody: WorkflowStateContractInput = {
+              ...workflowData,
+              loops: workflowData.loops || {},
+              parallels: workflowData.parallels || {},
+              edges: sanitizedEdges,
+            }
+
             try {
               await requestJson(putWorkflowNormalizedStateContract, {
                 params: { id: newWorkflow.id },
-                body: workflowData as ContractBodyInput<typeof putWorkflowNormalizedStateContract>,
+                body: workflowStateBody,
               })
             } catch (error) {
               logger.error(`Failed to save workflow state for ${newWorkflow.id}`, { error })

apps/sim/ee/audit-logs/hooks/audit-logs.ts

@@ -1,11 +1,6 @@
 import { keepPreviousData, useInfiniteQuery } from '@tanstack/react-query'
-import type { z } from 'zod'
 import { requestJson } from '@/lib/api/client/request'
-import type { ContractJsonResponse } from '@/lib/api/contracts'
-import {
-  type enterpriseAuditLogEntrySchema,
-  listAuditLogsContract,
-} from '@/lib/api/contracts/audit-logs'
+import { type AuditLogPage, listAuditLogsContract } from '@/lib/api/contracts/audit-logs'
 
 export const auditLogKeys = {
   all: ['audit-logs'] as const,
@@ -22,9 +17,6 @@ export interface AuditLogFilters {
   endDate?: string
 }
 
-export type EnterpriseAuditLogEntry = z.output<typeof enterpriseAuditLogEntrySchema>
-type AuditLogPage = ContractJsonResponse<typeof listAuditLogsContract>
-
 async function fetchAuditLogs(
   filters: AuditLogFilters,
   cursor?: string,

apps/sim/lib/api/contracts/audit-logs.ts

@@ -99,12 +99,16 @@ export const enterpriseAuditLogEntrySchema = z.object({
   createdAt: z.string(),
 })
 
+export type EnterpriseAuditLogEntry = z.output<typeof enterpriseAuditLogEntrySchema>
+
 export const listAuditLogsResponseSchema = z.object({
   success: z.boolean(),
   data: z.array(enterpriseAuditLogEntrySchema),
   nextCursor: z.string().optional(),
 })
 
+export type AuditLogPage = z.output<typeof listAuditLogsResponseSchema>
+
 export const listAuditLogsContract = defineRouteContract({
   method: 'GET',
   path: '/api/audit-logs',

apps/sim/lib/knowledge/documents/service.ts

@@ -47,6 +47,7 @@ import type { ProcessedDocumentTags } from '@/lib/knowledge/types'
 import { estimateTokenCount } from '@/lib/tokenization/estimators'
 import { deleteFile } from '@/lib/uploads/core/storage-service'
 import { extractStorageKey } from '@/lib/uploads/utils/file-utils'
+import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
 import type { DocumentProcessingPayload } from '@/background/knowledge-processing'
 import { calculateCost } from '@/providers/utils'
 
@@ -433,6 +434,13 @@ export async function processDocumentAsync(
     }
 
     const kbEmbeddingModel = kb[0].embeddingModel
+    if (!kb[0].workspaceId) {
+      throw new Error(`Knowledge base ${knowledgeBaseId} is missing workspace billing context`)
+    }
+    const billingUserId = await getWorkspaceBilledAccountUserId(kb[0].workspaceId)
+    if (!billingUserId) {
+      throw new Error(`Workspace ${kb[0].workspaceId} is missing billed account`)
+    }
     let totalEmbeddingTokens = 0
     let embeddingIsBYOK = false
     let embeddingModelName = kbEmbeddingModel
@@ -625,7 +633,7 @@ export async function processDocumentAsync(
     const processingTime = Date.now() - startTime
     logger.info(`[${documentId}] Successfully processed document in ${processingTime}ms`)
 
-    if (!embeddingIsBYOK && totalEmbeddingTokens > 0 && kb[0].userId) {
+    if (!embeddingIsBYOK && totalEmbeddingTokens > 0 && billingUserId) {
       try {
         const costMultiplier = getCostMultiplier()
         const { total: cost } = calculateCost(
@@ -637,7 +645,7 @@ export async function processDocumentAsync(
         )
         if (cost > 0) {
           await recordUsage({
-            userId: kb[0].userId,
+            userId: billingUserId,
             workspaceId: kb[0].workspaceId ?? undefined,
             entries: [
               {
@@ -652,7 +660,7 @@ export async function processDocumentAsync(
               totalTokensUsed: sql`total_tokens_used + ${totalEmbeddingTokens}`,
             },
           })
-          await checkAndBillOverageThreshold(kb[0].userId)
+          await checkAndBillOverageThreshold(billingUserId)
         } else {
           logger.warn(
             `[${documentId}] Embedding model "${embeddingModelName}" has no pricing entry — billing skipped`,