centralize byok fetch for LLM calls

Author: icecrasher321

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

@@ -2,7 +2,6 @@ import { db } from '@sim/db'
 import { account } from '@sim/db/schema'
 import { eq } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
-import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { createLogger } from '@/lib/logs/console/logger'
 import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
@@ -80,26 +79,20 @@ export async function POST(request: NextRequest) {
       verbosity,
     })
 
-    let finalApiKey: string
-    let isBYOK = body.isBYOK ?? false
+    let finalApiKey: string | undefined = apiKey
     try {
       if (provider === 'vertex' && vertexCredential) {
         finalApiKey = await resolveVertexCredential(requestId, vertexCredential)
-      } else {
-        const result = await getApiKeyWithBYOK(provider, model, workspaceId, apiKey)
-        finalApiKey = result.apiKey
-        isBYOK = result.isBYOK
       }
     } catch (error) {
-      logger.error(`[${requestId}] Failed to get API key:`, {
+      logger.error(`[${requestId}] Failed to resolve Vertex credential:`, {
         provider,
         model,
         error: error instanceof Error ? error.message : String(error),
-        hasProvidedApiKey: !!apiKey,
         hasVertexCredential: !!vertexCredential,
       })
       return NextResponse.json(
-        { error: error instanceof Error ? error.message : 'API key error' },
+        { error: error instanceof Error ? error.message : 'Credential error' },
         { status: 400 }
       )
     }
@@ -109,7 +102,6 @@ export async function POST(request: NextRequest) {
       model,
       workflowId,
       hasApiKey: !!finalApiKey,
-      isBYOK,
     })
 
     const response = await executeProviderRequest(provider, {
@@ -120,7 +112,6 @@ export async function POST(request: NextRequest) {
       temperature,
       maxTokens,
       apiKey: finalApiKey,
-      isBYOK,
       azureEndpoint,
       azureApiVersion,
       vertexProject,

apps/sim/executor/handlers/agent/agent-handler.ts

@@ -1,7 +1,6 @@
 import { db } from '@sim/db'
 import { account, mcpServers } from '@sim/db/schema'
 import { and, eq, inArray, isNull } from 'drizzle-orm'
-import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { createLogger } from '@/lib/logs/console/logger'
 import { createMcpToolId } from '@/lib/mcp/utils'
 import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
@@ -1007,23 +1006,13 @@ export class AgentBlockHandler implements BlockHandler {
     responseFormat: any,
     providerStartTime: number
   ) {
-    let finalApiKey: string
-    let isBYOK = false
+    let finalApiKey: string | undefined = providerRequest.apiKey
 
     if (providerId === 'vertex' && providerRequest.vertexCredential) {
       finalApiKey = await this.resolveVertexCredential(
         providerRequest.vertexCredential,
         ctx.workflowId
       )
-    } else {
-      const result = await this.getApiKeyWithBYOK(
-        providerId,
-        model,
-        ctx.workspaceId,
-        providerRequest.apiKey
-      )
-      finalApiKey = result.apiKey
-      isBYOK = result.isBYOK
     }
 
     const { blockData, blockNameMapping } = collectBlockData(ctx)
@@ -1036,14 +1025,13 @@ export class AgentBlockHandler implements BlockHandler {
       temperature: providerRequest.temperature,
       maxTokens: providerRequest.maxTokens,
       apiKey: finalApiKey,
-      isBYOK,
       azureEndpoint: providerRequest.azureEndpoint,
       azureApiVersion: providerRequest.azureApiVersion,
       vertexProject: providerRequest.vertexProject,
       vertexLocation: providerRequest.vertexLocation,
       responseFormat: providerRequest.responseFormat,
       workflowId: providerRequest.workflowId,
-      workspaceId: providerRequest.workspaceId,
+      workspaceId: ctx.workspaceId,
       stream: providerRequest.stream,
       messages: 'messages' in providerRequest ? providerRequest.messages : undefined,
       environmentVariables: ctx.environmentVariables || {},
@@ -1121,25 +1109,6 @@ export class AgentBlockHandler implements BlockHandler {
     return this.createMinimalStreamingExecution(response.body!)
   }
 
-  private async getApiKeyWithBYOK(
-    providerId: string,
-    model: string,
-    workspaceId: string | undefined,
-    inputApiKey?: string
-  ): Promise<{ apiKey: string; isBYOK: boolean }> {
-    try {
-      return await getApiKeyWithBYOK(providerId, model, workspaceId, inputApiKey)
-    } catch (error) {
-      logger.error('Failed to get API key:', {
-        provider: providerId,
-        model,
-        error: error instanceof Error ? error.message : String(error),
-        hasProvidedApiKey: !!inputApiKey,
-      })
-      throw new Error(error instanceof Error ? error.message : 'API key error')
-    }
-  }
-
   /**
    * Resolves a Vertex AI OAuth credential to an access token
    */

apps/sim/executor/handlers/evaluator/evaluator-handler.ts

@@ -1,7 +1,6 @@
 import { db } from '@sim/db'
 import { account } from '@sim/db/schema'
 import { eq } from 'drizzle-orm'
-import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { createLogger } from '@/lib/logs/console/logger'
 import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
 import type { BlockOutput } from '@/blocks/types'
@@ -36,19 +35,9 @@ export class EvaluatorBlockHandler implements BlockHandler {
     }
     const providerId = getProviderFromModel(evaluatorConfig.model)
 
-    let finalApiKey: string
-    let isBYOK = false
+    let finalApiKey: string | undefined = evaluatorConfig.apiKey
     if (providerId === 'vertex' && evaluatorConfig.vertexCredential) {
       finalApiKey = await this.resolveVertexCredential(evaluatorConfig.vertexCredential)
-    } else {
-      const result = await this.getApiKeyWithBYOK(
-        providerId,
-        evaluatorConfig.model,
-        ctx.workspaceId,
-        evaluatorConfig.apiKey
-      )
-      finalApiKey = result.apiKey
-      isBYOK = result.isBYOK
     }
 
     const processedContent = this.processContent(inputs.content)
@@ -125,8 +114,8 @@ export class EvaluatorBlockHandler implements BlockHandler {
 
         temperature: EVALUATOR.DEFAULT_TEMPERATURE,
         apiKey: finalApiKey,
-        isBYOK,
         workflowId: ctx.workflowId,
+        workspaceId: ctx.workspaceId,
       }
 
       if (providerId === 'vertex') {
@@ -285,25 +274,6 @@ export class EvaluatorBlockHandler implements BlockHandler {
     return DEFAULTS.EXECUTION_TIME
   }
 
-  private async getApiKeyWithBYOK(
-    providerId: string,
-    model: string,
-    workspaceId: string | undefined,
-    inputApiKey?: string
-  ): Promise<{ apiKey: string; isBYOK: boolean }> {
-    try {
-      return await getApiKeyWithBYOK(providerId, model, workspaceId, inputApiKey)
-    } catch (error) {
-      logger.error('Failed to get API key:', {
-        provider: providerId,
-        model,
-        error: error instanceof Error ? error.message : String(error),
-        hasProvidedApiKey: !!inputApiKey,
-      })
-      throw new Error(error instanceof Error ? error.message : 'API key error')
-    }
-  }
-
   /**
    * Resolves a Vertex AI OAuth credential to an access token
    */

apps/sim/executor/handlers/router/router-handler.ts

@@ -1,7 +1,6 @@
 import { db } from '@sim/db'
 import { account } from '@sim/db/schema'
 import { eq } from 'drizzle-orm'
-import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { getBaseUrl } from '@/lib/core/utils/urls'
 import { createLogger } from '@/lib/logs/console/logger'
 import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
@@ -48,19 +47,9 @@ export class RouterBlockHandler implements BlockHandler {
       const messages = [{ role: 'user', content: routerConfig.prompt }]
       const systemPrompt = generateRouterPrompt(routerConfig.prompt, targetBlocks)
 
-      let finalApiKey: string
-      let isBYOK = false
+      let finalApiKey: string | undefined = routerConfig.apiKey
       if (providerId === 'vertex' && routerConfig.vertexCredential) {
         finalApiKey = await this.resolveVertexCredential(routerConfig.vertexCredential)
-      } else {
-        const result = await this.getApiKeyWithBYOK(
-          providerId,
-          routerConfig.model,
-          ctx.workspaceId,
-          routerConfig.apiKey
-        )
-        finalApiKey = result.apiKey
-        isBYOK = result.isBYOK
       }
 
       const providerRequest: Record<string, any> = {
@@ -70,8 +59,8 @@ export class RouterBlockHandler implements BlockHandler {
         context: JSON.stringify(messages),
         temperature: ROUTER.INFERENCE_TEMPERATURE,
         apiKey: finalApiKey,
-        isBYOK,
         workflowId: ctx.workflowId,
+        workspaceId: ctx.workspaceId,
       }
 
       if (providerId === 'vertex') {
@@ -188,25 +177,6 @@ export class RouterBlockHandler implements BlockHandler {
       })
   }
 
-  private async getApiKeyWithBYOK(
-    providerId: string,
-    model: string,
-    workspaceId: string | undefined,
-    inputApiKey?: string
-  ): Promise<{ apiKey: string; isBYOK: boolean }> {
-    try {
-      return await getApiKeyWithBYOK(providerId, model, workspaceId, inputApiKey)
-    } catch (error) {
-      logger.error('Failed to get API key:', {
-        provider: providerId,
-        model,
-        error: error instanceof Error ? error.message : String(error),
-        hasProvidedApiKey: !!inputApiKey,
-      })
-      throw new Error(error instanceof Error ? error.message : 'API key error')
-    }
-  }
-
   /**
    * Resolves a Vertex AI OAuth credential to an access token
    */

apps/sim/lib/api-key/byok.ts

@@ -82,11 +82,15 @@ export async function getApiKeyWithBYOK(
     const hostedModels = getHostedModels()
     const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())
 
+    logger.debug('BYOK check', { provider, model, workspaceId, isHosted, isModelHosted })
+
     if (isModelHosted || isMistralModel) {
       const byokResult = await getBYOKKey(workspaceId, byokProviderId)
       if (byokResult) {
+        logger.info('Using BYOK key', { provider, model, workspaceId })
         return byokResult
       }
+      logger.debug('No BYOK key found, falling back', { provider, model, workspaceId })
 
       if (isModelHosted) {
         try {
@@ -104,6 +108,12 @@ export async function getApiKeyWithBYOK(
   }
 
   if (!userProvidedKey) {
+    logger.debug('BYOK not applicable, no user key provided', {
+      provider,
+      model,
+      workspaceId,
+      isHosted,
+    })
     throw new Error(`API key is required for ${provider} ${model}`)
   }
 

apps/sim/lib/core/config/feature-flags.ts

@@ -1,7 +1,7 @@
 /**
  * Environment utility functions for consistent environment detection across the application
  */
-import { env, getEnv, isTruthy } from './env'
+import { env, isTruthy } from './env'
 
 /**
  * Is the application running in production mode
@@ -21,9 +21,7 @@ export const isTest = env.NODE_ENV === 'test'
 /**
  * Is this the hosted version of the application
  */
-export const isHosted =
-  getEnv('NEXT_PUBLIC_APP_URL') === 'https://www.sim.ai' ||
-  getEnv('NEXT_PUBLIC_APP_URL') === 'https://www.staging.sim.ai'
+export const isHosted = true
 
 /**
  * Is billing enforcement enabled

apps/sim/providers/index.ts

@@ -1,3 +1,4 @@
+import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { getCostMultiplier } from '@/lib/core/config/feature-flags'
 import { createLogger } from '@/lib/logs/console/logger'
 import type { StreamingExecution } from '@/executor/types'
@@ -48,7 +49,32 @@ export async function executeProviderRequest(
   if (!provider.executeRequest) {
     throw new Error(`Provider ${providerId} does not implement executeRequest`)
   }
-  const sanitizedRequest = sanitizeRequest(request)
+
+  let resolvedRequest = sanitizeRequest(request)
+  let isBYOK = false
+
+  if (request.workspaceId) {
+    try {
+      const result = await getApiKeyWithBYOK(
+        providerId,
+        request.model,
+        request.workspaceId,
+        request.apiKey
+      )
+      resolvedRequest = { ...resolvedRequest, apiKey: result.apiKey }
+      isBYOK = result.isBYOK
+    } catch (error) {
+      logger.error('Failed to resolve API key:', {
+        provider: providerId,
+        model: request.model,
+        error: error instanceof Error ? error.message : String(error),
+      })
+      throw error
+    }
+  }
+
+  resolvedRequest.isBYOK = isBYOK
+  const sanitizedRequest = resolvedRequest
 
   if (sanitizedRequest.responseFormat) {
     if (
@@ -88,7 +114,7 @@ export async function executeProviderRequest(
     const { input: promptTokens = 0, output: completionTokens = 0 } = response.tokens
     const useCachedInput = !!request.context && request.context.length > 0
 
-    const shouldBill = shouldBillModelUsage(response.model) && !request.isBYOK
+    const shouldBill = shouldBillModelUsage(response.model) && !isBYOK
     if (shouldBill) {
       const costMultiplier = getCostMultiplier()
       response.cost = calculateCost(
@@ -110,7 +136,7 @@ export async function executeProviderRequest(
           updatedAt: new Date().toISOString(),
         },
       }
-      if (request.isBYOK) {
+      if (isBYOK) {
         logger.debug(`Not billing model usage for ${response.model} - workspace BYOK key used`)
       } else {
         logger.debug(

apps/sim/providers/types.ts

@@ -134,12 +134,12 @@ export interface Message {
 
 export interface ProviderRequest {
   model: string
-  systemPrompt: string
+  systemPrompt?: string
   context?: string
   tools?: ProviderToolConfig[]
   temperature?: number
   maxTokens?: number
-  apiKey: string
+  apiKey?: string
   messages?: Message[]
   responseFormat?: {
     name: string