feat(copilot): send workspace entitlements to the mothership

Author: TheodoreSpeaks

apps/sim/app/api/mothership/execute/route.ts

@@ -18,6 +18,7 @@ import type { StreamEvent } from '@/lib/copilot/request/types'
 import { isE2BDocEnabled } from '@/lib/core/config/env-flags'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { buildUserSkillTool } from '@/lib/mothership/skills'
+import { computeWorkspaceEntitlements } from '@/lib/workflows/custom-blocks/operations'
 import {
   assertActiveWorkspaceAccess,
   getUserEntityPermissions,
@@ -141,25 +142,32 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
     const lastUserMessage = messages.filter((m) => m.role === 'user').at(-1)?.content
     // double-cast-allowed: the contract validates contexts as open kind/label objects; processContextsServer narrows on `kind` at runtime
     const agentMentions = contexts as unknown as ChatContext[] | undefined
-    const [workspaceContext, integrationTools, userSkillTool, userPermission, agentContexts] =
-      await Promise.all([
-        generateWorkspaceContext(workspaceId, userId),
-        buildIntegrationToolSchemas(userId, messageId, undefined, workspaceId),
-        buildUserSkillTool(workspaceId),
-        getUserEntityPermissions(userId, 'workspace', workspaceId).catch(() => null),
-        processContextsServer(
-          agentMentions,
-          userId,
-          lastUserMessage,
-          workspaceId,
-          effectiveChatId
-        ).catch((error) => {
-          reqLogger.warn('Failed to resolve agent contexts for execution', {
-            error: toError(error).message,
-          })
-          return []
-        }),
-      ])
+    const [
+      workspaceContext,
+      integrationTools,
+      userSkillTool,
+      userPermission,
+      entitlements,
+      agentContexts,
+    ] = await Promise.all([
+      generateWorkspaceContext(workspaceId, userId),
+      buildIntegrationToolSchemas(userId, messageId, undefined, workspaceId),
+      buildUserSkillTool(workspaceId),
+      getUserEntityPermissions(userId, 'workspace', workspaceId).catch(() => null),
+      computeWorkspaceEntitlements(workspaceId, userId),
+      processContextsServer(
+        agentMentions,
+        userId,
+        lastUserMessage,
+        workspaceId,
+        effectiveChatId
+      ).catch((error) => {
+        reqLogger.warn('Failed to resolve agent contexts for execution', {
+          error: toError(error).message,
+        })
+        return []
+      }),
+    ])
     const requestPayload: Record<string, unknown> = {
       messages,
       responseFormat,
@@ -182,6 +190,7 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
       ...(integrationTools.length > 0 ? { integrationTools } : {}),
       ...(userSkillTool ? { mothershipTools: [userSkillTool] } : {}),
       ...(userPermission ? { userPermission } : {}),
+      ...(entitlements.length > 0 ? { entitlements } : {}),
     }
 
     let allowExplicitAbort = true

apps/sim/lib/copilot/chat/payload.test.ts

@@ -227,4 +227,34 @@ describe('buildCopilotRequestPayload', () => {
       })
     )
   })
+
+  it('passes entitlements through and omits the field when empty', async () => {
+    const withEntitlements = await buildCopilotRequestPayload(
+      {
+        message: 'publish as a block',
+        userId: 'user-1',
+        userMessageId: 'msg-1',
+        mode: 'agent',
+        model: 'claude-opus-4-8',
+        workspaceId: 'ws-1',
+        entitlements: ['custom-blocks'],
+      },
+      { selectedModel: 'claude-opus-4-8' }
+    )
+    expect(withEntitlements).toEqual(expect.objectContaining({ entitlements: ['custom-blocks'] }))
+
+    const withoutEntitlements = await buildCopilotRequestPayload(
+      {
+        message: 'publish as a block',
+        userId: 'user-1',
+        userMessageId: 'msg-1',
+        mode: 'agent',
+        model: 'claude-opus-4-8',
+        workspaceId: 'ws-1',
+        entitlements: [],
+      },
+      { selectedModel: 'claude-opus-4-8' }
+    )
+    expect(withoutEntitlements).not.toHaveProperty('entitlements')
+  })
 })

apps/sim/lib/copilot/chat/payload.ts

@@ -36,6 +36,8 @@ interface BuildPayloadParams {
   workspaceContext?: string
   vfs?: VfsSnapshotV1
   userPermission?: string
+  /** Plan/flag-gated org capabilities (e.g. "custom-blocks") the mothership gates tools/prompts on. */
+  entitlements?: string[]
   userTimezone?: string
   userMetadata?: {
     name?: string
@@ -370,6 +372,7 @@ export async function buildCopilotRequestPayload(
     ...(params.workspaceContext ? { workspaceContext: params.workspaceContext } : {}),
     ...(params.vfs ? { vfs: params.vfs } : {}),
     ...(params.userPermission ? { userPermission: params.userPermission } : {}),
+    ...(params.entitlements?.length ? { entitlements: params.entitlements } : {}),
     ...(params.userTimezone ? { userTimezone: params.userTimezone } : {}),
     ...(params.userMetadata &&
     (params.userMetadata.name || params.userMetadata.email || params.userMetadata.timezone)

apps/sim/lib/copilot/chat/post.ts

@@ -45,6 +45,7 @@ import type { ExecutionContext, OrchestratorResult } from '@/lib/copilot/request
 import { persistChatResources } from '@/lib/copilot/resources/persistence'
 import { prepareExecutionContext } from '@/lib/copilot/tools/handlers/context'
 import { getEffectiveDecryptedEnv } from '@/lib/environment/utils'
+import { computeWorkspaceEntitlements } from '@/lib/workflows/custom-blocks/operations'
 import { resolveWorkflowIdForUser } from '@/lib/workflows/utils'
 import {
   getUserEntityPermissions,
@@ -174,6 +175,7 @@ type UnifiedChatBranch =
         contexts: Array<{ type: string; content: string; tag?: string; path?: string }>
         fileAttachments?: UnifiedChatRequest['fileAttachments']
         userPermission?: string
+        entitlements?: string[]
         userTimezone?: string
         userMetadata?: { name?: string; email?: string; timezone?: string }
         workflowId: string
@@ -211,6 +213,7 @@ type UnifiedChatBranch =
         contexts: Array<{ type: string; content: string; tag?: string; path?: string }>
         fileAttachments?: UnifiedChatRequest['fileAttachments']
         userPermission?: string
+        entitlements?: string[]
         userTimezone?: string
         userMetadata?: { name?: string; email?: string; timezone?: string }
         workspaceContext?: string
@@ -623,6 +626,7 @@ async function resolveBranch(params: {
             workspaceContext: payloadParams.workspaceContext,
             vfs: payloadParams.vfs,
             userPermission: payloadParams.userPermission,
+            entitlements: payloadParams.entitlements,
             userTimezone: payloadParams.userTimezone,
             userMetadata: payloadParams.userMetadata,
           },
@@ -678,6 +682,7 @@ async function resolveBranch(params: {
           workspaceContext: payloadParams.workspaceContext,
           vfs: payloadParams.vfs,
           userPermission: payloadParams.userPermission,
+          entitlements: payloadParams.entitlements,
           userTimezone: payloadParams.userTimezone,
           userMetadata: payloadParams.userMetadata,
           includeMothershipTools: true,
@@ -898,6 +903,9 @@ export async function handleUnifiedChatPost(req: NextRequest) {
                 }
               )
             : Promise.resolve(null)
+      const entitlementsPromise = workspaceId
+        ? computeWorkspaceEntitlements(workspaceId, authenticatedUserId)
+        : Promise.resolve([])
       // Wrap the pre-LLM prep work in spans so the trace waterfall shows
       // where time is going between "request received" and "llm.stream
       // opens". Previously these ran bare under the root and inflated the
@@ -952,10 +960,11 @@ export async function handleUnifiedChatPost(req: NextRequest) {
         activeOtelRoot.context
       )
 
-      const [agentContexts, userPermission, workspaceSnapshot, , executionContext] =
+      const [agentContexts, userPermission, entitlements, workspaceSnapshot, , executionContext] =
         await Promise.all([
           agentContextsPromise,
           userPermissionPromise,
+          entitlementsPromise,
           workspaceContextPromise,
           persistUserMessagePromise,
           executionContextPromise,
@@ -990,6 +999,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
                 contexts: agentContexts,
                 fileAttachments: body.fileAttachments,
                 userPermission: userPermission ?? undefined,
+                entitlements,
                 userTimezone: body.userTimezone,
                 userMetadata,
                 workflowId: branch.workflowId,
@@ -1011,6 +1021,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
                 contexts: agentContexts,
                 fileAttachments: body.fileAttachments,
                 userPermission: userPermission ?? undefined,
+                entitlements,
                 userTimezone: body.userTimezone,
                 userMetadata,
                 workspaceContext,

apps/sim/lib/copilot/tools/handlers/deployment/custom-block.test.ts

@@ -142,6 +142,25 @@ describe('executeDeployCustomBlock', () => {
     })
   })
 
+  it('returns a clean admin-permission error when workflow access is denied', async () => {
+    ensureWorkflowAccessMock.mockRejectedValue(new Error('Unauthorized workflow access'))
+
+    const result = await executeDeployCustomBlock({ name: 'Enrich Lead' }, context)
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('admin permission')
+    expect(publishCustomBlockMock).not.toHaveBeenCalled()
+  })
+
+  it('surfaces workflow-not-found from access resolution', async () => {
+    ensureWorkflowAccessMock.mockRejectedValue(new Error('Workflow wf-1 not found'))
+
+    const result = await executeDeployCustomBlock({ name: 'Enrich Lead' }, context)
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('not found')
+  })
+
   it('requires a name on first publish', async () => {
     const result = await executeDeployCustomBlock({}, context)
 

apps/sim/lib/copilot/tools/handlers/deployment/custom-block.ts

@@ -117,11 +117,19 @@ export async function executeDeployCustomBlock(
     }
     const action = params.action === 'undeploy' ? 'undeploy' : 'deploy'
 
-    const { workflow: workflowRecord } = await ensureWorkflowAccess(
-      workflowId,
-      context.userId,
-      'admin'
-    )
+    let workflowRecord: Awaited<ReturnType<typeof ensureWorkflowAccess>>['workflow']
+    try {
+      workflowRecord = (await ensureWorkflowAccess(workflowId, context.userId, 'admin')).workflow
+    } catch (error) {
+      const message = toError(error).message
+      if (message.includes('not found')) {
+        return { success: false, error: message }
+      }
+      return {
+        success: false,
+        error: "Managing a custom block requires admin permission on the workflow's workspace",
+      }
+    }
     const workspaceId = workflowRecord.workspaceId
     if (!workspaceId) {
       return { success: false, error: 'Workflow must belong to a workspace' }

apps/sim/lib/mothership/inbox/executor.ts

@@ -24,6 +24,7 @@ import type { AgentMailAttachment } from '@/lib/mothership/inbox/types'
 import { buildUserSkillTool } from '@/lib/mothership/skills'
 import { uploadFile } from '@/lib/uploads/core/storage-service'
 import { createFileContent, type MessageContent } from '@/lib/uploads/utils/file-utils'
+import { computeWorkspaceEntitlements } from '@/lib/workflows/custom-blocks/operations'
 import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
 
@@ -208,14 +209,21 @@ export async function executeInboxTask(taskId: string): Promise<void> {
       return { attachments, ...downloaded }
     }
 
-    const [attachmentResult, workspaceContext, integrationTools, userSkillTool, userPermission] =
-      await Promise.all([
-        fetchAttachments(),
-        generateWorkspaceContext(ws.id, userId),
-        buildIntegrationToolSchemas(userId, undefined, undefined, ws.id),
-        buildUserSkillTool(ws.id),
-        getUserEntityPermissions(userId, 'workspace', ws.id).catch(() => null),
-      ])
+    const [
+      attachmentResult,
+      workspaceContext,
+      integrationTools,
+      userSkillTool,
+      userPermission,
+      entitlements,
+    ] = await Promise.all([
+      fetchAttachments(),
+      generateWorkspaceContext(ws.id, userId),
+      buildIntegrationToolSchemas(userId, undefined, undefined, ws.id),
+      buildUserSkillTool(ws.id),
+      getUserEntityPermissions(userId, 'workspace', ws.id).catch(() => null),
+      computeWorkspaceEntitlements(ws.id, userId),
+    ])
     const { attachments, fileAttachments, storedAttachments } = attachmentResult
 
     const truncatedTask = {
@@ -237,6 +245,7 @@ export async function executeInboxTask(taskId: string): Promise<void> {
       ...(integrationTools.length > 0 ? { integrationTools } : {}),
       ...(userSkillTool ? { mothershipTools: [userSkillTool] } : {}),
       ...(userPermission ? { userPermission } : {}),
+      ...(entitlements.length > 0 ? { entitlements } : {}),
       ...(fileAttachments.length > 0 ? { fileAttachments } : {}),
     }
 

apps/sim/lib/workflows/custom-blocks/operations.ts

@@ -1,8 +1,10 @@
 import { db } from '@sim/db'
 import { customBlock, workflow, workspace } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
 import { generateId, generateShortId } from '@sim/utils/id'
 import { and, eq } from 'drizzle-orm'
+import { LRUCache } from 'lru-cache'
 import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription'
 import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
 import { extractInputFieldsFromBlocks, type WorkflowInputField } from '@/lib/workflows/input-format'
@@ -19,15 +21,65 @@ const logger = createLogger('CustomBlocksOperations')
  * enterprise plan). Applying it in every org-scoped resolver keeps execution, the
  * copilot VFS, and workspace context from surfacing blocks the API withholds (e.g.
  * after an org drops off the enterprise plan). Returns `null` when ineligible.
+ *
+ * Pass `userId` when the caller acts for a specific user so per-user flag
+ * targeting matches the REST routes; workspace-scoped resolvers (VFS, context,
+ * executor overlay) omit it and evaluate at org level.
  */
-async function eligibleOrgForWorkspace(workspaceId: string): Promise<string | null> {
+async function eligibleOrgForWorkspace(
+  workspaceId: string,
+  userId?: string
+): Promise<string | null> {
   const ws = await getWorkspaceWithOwner(workspaceId, { includeArchived: true })
   if (!ws?.organizationId) return null
-  if (!(await isFeatureEnabled('deploy-as-block', { orgId: ws.organizationId }))) return null
+  if (!(await isFeatureEnabled('deploy-as-block', { userId, orgId: ws.organizationId }))) {
+    return null
+  }
   if (!(await isOrganizationOnEnterprisePlan(ws.organizationId))) return null
   return ws.organizationId
 }
 
+/**
+ * Cross-repo contract: the mothership (Go) reads this exact string from the chat
+ * payload's `entitlements` array and matches it against
+ * `core.EntitlementCustomBlocks` to gate the deploy agent's custom-block tool,
+ * prompt section, and skill. The entitlement is the materialized form of
+ * "`deploy-as-block` flag AND enterprise plan" for the workspace's org.
+ */
+export const CUSTOM_BLOCKS_ENTITLEMENT = 'custom-blocks'
+
+const entitlementsCache = new LRUCache<string, Promise<string[]>>({
+  max: 500,
+  ttl: 5_000,
+})
+
+/**
+ * The plan/flag-gated capabilities to send to the mothership for a request in
+ * this workspace. Cached briefly so the several per-message callers (payload
+ * build, headless execute, inbox) collapse to one evaluation.
+ */
+export function computeWorkspaceEntitlements(
+  workspaceId: string,
+  userId?: string
+): Promise<string[]> {
+  const cacheKey = `${workspaceId}:${userId ?? ''}`
+  const cached = entitlementsCache.get(cacheKey)
+  if (cached) return cached
+
+  const promise = eligibleOrgForWorkspace(workspaceId, userId)
+    .then((organizationId) => (organizationId ? [CUSTOM_BLOCKS_ENTITLEMENT] : []))
+    .catch((error) => {
+      entitlementsCache.delete(cacheKey)
+      logger.warn('Failed to compute workspace entitlements', {
+        workspaceId,
+        error: getErrorMessage(error),
+      })
+      return []
+    })
+  entitlementsCache.set(cacheKey, promise)
+  return promise
+}
+
 /** A persisted custom block plus its live-derived Start input fields. */
 export interface CustomBlockWithInputs {
   id: string