Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 28 additions & 19 deletions apps/sim/app/api/mothership/execute/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { checkInternalAuth } from '@/lib/auth/hybrid'
import { buildIntegrationToolSchemas } from '@/lib/copilot/chat/payload'
import { processContextsServer } from '@/lib/copilot/chat/process-contents'
import { generateWorkspaceContext } from '@/lib/copilot/chat/workspace-context'
import { computeWorkspaceEntitlements } from '@/lib/copilot/entitlements'
import {
MothershipStreamV1EventType,
MothershipStreamV1TextChannel,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
30 changes: 30 additions & 0 deletions apps/sim/lib/copilot/chat/payload.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -237,4 +237,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')
})
})
3 changes: 3 additions & 0 deletions apps/sim/lib/copilot/chat/payload.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,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
Expand Down Expand Up @@ -387,6 +389,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)
Expand Down
13 changes: 12 additions & 1 deletion apps/sim/lib/copilot/chat/post.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import { finalizeAssistantTurn } from '@/lib/copilot/chat/terminal-state'
import { generateWorkspaceSnapshot } from '@/lib/copilot/chat/workspace-context'
import { chatPubSub } from '@/lib/copilot/chat-status'
import { COPILOT_REQUEST_MODES } from '@/lib/copilot/constants'
import { computeWorkspaceEntitlements } from '@/lib/copilot/entitlements'
import {
CopilotChatFinalizeOutcome,
CopilotChatPersistOutcome,
Expand Down Expand Up @@ -175,6 +176,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
Expand Down Expand Up @@ -212,6 +214,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
Expand Down Expand Up @@ -624,6 +627,7 @@ async function resolveBranch(params: {
workspaceContext: payloadParams.workspaceContext,
vfs: payloadParams.vfs,
userPermission: payloadParams.userPermission,
entitlements: payloadParams.entitlements,
userTimezone: payloadParams.userTimezone,
userMetadata: payloadParams.userMetadata,
},
Expand Down Expand Up @@ -679,6 +683,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,
Expand Down Expand Up @@ -899,6 +904,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
Expand Down Expand Up @@ -953,10 +961,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,
Expand Down Expand Up @@ -991,6 +1000,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
contexts: agentContexts,
fileAttachments: body.fileAttachments,
userPermission: userPermission ?? undefined,
entitlements,
userTimezone: body.userTimezone,
userMetadata,
workflowId: branch.workflowId,
Expand All @@ -1012,6 +1022,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
contexts: agentContexts,
fileAttachments: body.fileAttachments,
userPermission: userPermission ?? undefined,
entitlements,
userTimezone: body.userTimezone,
userMetadata,
workspaceContext,
Expand Down
72 changes: 72 additions & 0 deletions apps/sim/lib/copilot/entitlements.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { LRUCache } from 'lru-cache'
import { isCustomBlocksEligible } from '@/lib/workflows/custom-blocks/operations'

const logger = createLogger('CopilotEntitlements')

/**
* Cross-repo contract: the mothership (Go) matches these exact strings against
* its `core.Entitlement*` constants to gate agent surfaces.
*/
export const CUSTOM_BLOCKS_ENTITLEMENT = 'custom-blocks'

/**
* Workspace entitlements — plan/flag-gated org capabilities sent to the
* mothership as the chat payload's `entitlements` array. The Go side hides the
* matching tools, skills, and prompt sections when an entitlement is absent, so
* a non-entitled org's agents never hear of the feature.
*
* Adding an entitlement:
* 1. Add the kebab-case name and a fail-closed evaluator here. Every payload
* site (interactive chat, headless execute, inbox) picks it up automatically.
* 2. Go repo: add the matching `Entitlement*` constant in `internal/core` and
* gate surfaces declaratively — `RequiredEntitlement` on tool definitions,
* `entitlement:` frontmatter on skills, a conditional section in
* `BuildAgentEnvelope`, or a variant swap in `Capabilities`.
* 3. Keep enforcement in sim: the Go gating is advertisement-only (the payload
* is forgeable), so the sim-side tool handler must re-check the same
* predicate at execution time.
*/
const ENTITLEMENT_EVALUATORS: Record<
string,
(workspaceId: string, userId?: string) => Promise<boolean>
> = {
[CUSTOM_BLOCKS_ENTITLEMENT]: isCustomBlocksEligible,
}

const entitlementsCache = new LRUCache<string, Promise<string[]>>({
max: 500,
ttl: 5_000,
})

/**
* The entitlements to send to the mothership for a request in this workspace.
* Each evaluator fails closed (an error means the entitlement is absent).
* Cached briefly so the several per-message callers 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 = Promise.all(
Object.entries(ENTITLEMENT_EVALUATORS).map(async ([name, evaluate]) => {
try {
return (await evaluate(workspaceId, userId)) ? name : null
} catch (error) {
logger.warn('Entitlement evaluation failed; treating as absent', {
entitlement: name,
workspaceId,
error: getErrorMessage(error),
})
return null
}
})
).then((names) => names.filter((name): name is string => name !== null))
entitlementsCache.set(cacheKey, promise)
return promise
}
Loading
Loading