chore: import upstream snapshot with attribution

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

@@ -0,0 +1,556 @@
+import { db } from '@sim/db'
+import {
+  customBlock,
+  workflow,
+  workflowBlocks,
+  workflowDeploymentVersion,
+  workspace,
+} from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { generateId, generateShortId } from '@sim/utils/id'
+import { and, eq, isNull, sql } from 'drizzle-orm'
+import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription'
+import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
+import { extractInputFieldsFromBlocks, type WorkflowInputField } from '@/lib/workflows/input-format'
+import { loadDeployedWorkflowState } from '@/lib/workflows/persistence/utils'
+import { getWorkspaceWithOwner } from '@/lib/workspaces/permissions/utils'
+import type { CustomBlockOutput, CustomBlockRow } from '@/blocks/custom/build-config'
+import { CUSTOM_BLOCK_TYPE_PREFIX } from '@/blocks/custom/build-config'
+
+const logger = createLogger('CustomBlocksOperations')
+
+/**
+ * Resolve a workspace's organization ONLY when custom blocks are enabled for it —
+ * the same gate the REST list/publish routes apply (`deploy-as-block` flag +
+ * 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,
+  userId?: string
+): Promise<string | null> {
+  const ws = await getWorkspaceWithOwner(workspaceId, { includeArchived: true })
+  if (!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
+}
+
+/**
+ * Whether the workspace's org may use custom blocks (`deploy-as-block` flag +
+ * enterprise plan). Feeds the 'custom-blocks' entitlement in
+ * `@/lib/copilot/entitlements` and matches the REST route gates.
+ */
+export async function isCustomBlocksEligible(
+  workspaceId: string,
+  userId?: string
+): Promise<boolean> {
+  return (await eligibleOrgForWorkspace(workspaceId, userId)) !== null
+}
+
+/** A persisted custom block plus its live-derived Start input fields. */
+export interface CustomBlockWithInputs {
+  id: string
+  organizationId: string
+  workflowId: string
+  workflowName: string
+  /** Source workflow's home workspace id — used client-side to gate manage affordances. */
+  workspaceId: string | null
+  workspaceName: string | null
+  type: string
+  name: string
+  description: string
+  iconUrl: string | null
+  enabled: boolean
+  inputFields: WorkflowInputField[]
+  exposedOutputs: CustomBlockOutput[]
+}
+
+/**
+ * Derive a bound workflow's Start input fields from its LATEST DEPLOYMENT — the
+ * exact state execution runs. Deriving from the draft/editor tables would let the
+ * block advertise inputs the deployed child doesn't accept (or miss ones it still
+ * expects) whenever the publisher edits after deploying. Returns `[]` if the
+ * workflow has no active deployment.
+ */
+async function deriveInputFields(workflowId: string): Promise<WorkflowInputField[]> {
+  try {
+    const deployed = await loadDeployedWorkflowState(workflowId)
+    return extractInputFieldsFromBlocks(deployed.blocks)
+  } catch {
+    return []
+  }
+}
+
+/** A stored per-input override (placeholder + required), keyed by the Start field's stable id. */
+type InputPlaceholder = { id: string; placeholder?: string; required?: boolean }
+
+/**
+ * The block's input fields: the LIVE deployed Start fields (authoritative for which
+ * inputs exist and their name/type — so an input removed from the source and
+ * redeployed simply disappears) with the stored per-id `placeholder`/`required`
+ * overrides merged in. When the source is undeployed there are no live fields, so
+ * there are no inputs — the block can't run undeployed anyway.
+ */
+function applyInputPlaceholders(
+  placeholders: InputPlaceholder[] | null,
+  deployed: WorkflowInputField[]
+): WorkflowInputField[] {
+  if (deployed.length === 0) return []
+  if (!placeholders?.length) return deployed
+  const byId = new Map(placeholders.map((p) => [p.id, p]))
+  // Overrides are stored under `field.id ?? field.name` (the form's key), so a
+  // legacy field with no stable id is keyed by name — look it up the same way.
+  return deployed.map((field) => {
+    const override = byId.get(field.id ?? field.name)
+    if (!override) return field
+    return {
+      ...field,
+      ...(override.placeholder ? { placeholder: override.placeholder } : {}),
+      ...(override.required ? { required: true } : {}),
+    }
+  })
+}
+
+/**
+ * The org's custom blocks for the server overlay (`withCustomBlockOverlay`).
+ * Includes DISABLED rows (carrying `enabled`) so a still-placed disabled block
+ * stays resolvable — it survives serialization and fails loudly at run via
+ * `getCustomBlockAuthority` instead of being silently dropped from the graph; the
+ * overlay marks it `hideFromToolbar` so no new instance can be placed. No input
+ * fields: the server's `inputMapping` is schema-agnostic and the handler's remap
+ * filters every value against the child's live deployed Start.
+ */
+export async function getCustomBlockRowsForOrg(
+  organizationId: string
+): Promise<Array<CustomBlockRow & { enabled: boolean }>> {
+  const rows = await db
+    .select({
+      type: customBlock.type,
+      name: customBlock.name,
+      description: customBlock.description,
+      workflowId: customBlock.workflowId,
+      outputs: customBlock.outputs,
+      enabled: customBlock.enabled,
+    })
+    .from(customBlock)
+    .where(eq(customBlock.organizationId, organizationId))
+
+  return rows.map(({ outputs, ...r }) => ({ ...r, exposedOutputs: outputs ?? [] }))
+}
+
+/**
+ * The custom-block rows in scope for a workspace's organization, for wrapping an
+ * execution in `withCustomBlockOverlay`. Returns `[]` when the workspace has no
+ * organization (nothing to resolve).
+ */
+export async function getCustomBlockRowsForWorkspace(
+  workspaceId: string
+): Promise<CustomBlockRow[]> {
+  const organizationId = await eligibleOrgForWorkspace(workspaceId)
+  if (!organizationId) return []
+  return getCustomBlockRowsForOrg(organizationId)
+}
+
+/**
+ * The custom blocks (with live-derived input fields) for a workspace's org. Used
+ * by the copilot VFS to expose custom blocks to the agent. Returns `[]` when the
+ * workspace has no organization.
+ */
+export async function listCustomBlocksWithInputsForWorkspace(
+  workspaceId: string
+): Promise<CustomBlockWithInputs[]> {
+  const organizationId = await eligibleOrgForWorkspace(workspaceId)
+  if (!organizationId) return []
+  return listCustomBlocksWithInputs(organizationId)
+}
+
+/**
+ * Lightweight enabled-custom-block summaries for a workspace's org (type + name +
+ * description, no input derivation). Used by the copilot workspace-context markdown.
+ */
+export async function listCustomBlockSummariesForWorkspace(
+  workspaceId: string
+): Promise<Array<{ type: string; name: string; description: string }>> {
+  const organizationId = await eligibleOrgForWorkspace(workspaceId)
+  if (!organizationId) return []
+  return db
+    .select({
+      type: customBlock.type,
+      name: customBlock.name,
+      description: customBlock.description,
+    })
+    .from(customBlock)
+    .where(and(eq(customBlock.organizationId, organizationId), eq(customBlock.enabled, true)))
+}
+
+/**
+ * Hydrate a joined custom-block row into the wire shape. Field set derived live
+ * from the deployed Start; stored placeholders merged in. Derive even for a
+ * disabled block — the source workflow's deployment is independent of the block's
+ * enabled flag, and the edit form needs the real fields so a save doesn't
+ * overwrite the block's stored placeholders.
+ */
+async function hydrateCustomBlockRow(joined: {
+  block: typeof customBlock.$inferSelect
+  workflowName: string
+  workspaceId: string | null
+  workspaceName: string | null
+}): Promise<CustomBlockWithInputs> {
+  const { block: row, workflowName, workspaceId, workspaceName } = joined
+  return {
+    id: row.id,
+    organizationId: row.organizationId,
+    workflowId: row.workflowId,
+    workflowName,
+    workspaceId,
+    workspaceName,
+    type: row.type,
+    name: row.name,
+    description: row.description,
+    iconUrl: row.iconUrl,
+    enabled: row.enabled,
+    inputFields: applyInputPlaceholders(row.inputs, await deriveInputFields(row.workflowId)),
+    exposedOutputs: row.outputs ?? [],
+  }
+}
+
+/** The org's custom blocks with live-derived input fields (client overlay + list API). */
+export async function listCustomBlocksWithInputs(
+  organizationId: string
+): Promise<CustomBlockWithInputs[]> {
+  const rows = await db
+    .select({
+      block: customBlock,
+      workflowName: workflow.name,
+      workspaceId: workflow.workspaceId,
+      workspaceName: workspace.name,
+    })
+    .from(customBlock)
+    .innerJoin(workflow, eq(workflow.id, customBlock.workflowId))
+    .leftJoin(workspace, eq(workspace.id, workflow.workspaceId))
+    .where(eq(customBlock.organizationId, organizationId))
+
+  return Promise.all(rows.map(hydrateCustomBlockRow))
+}
+
+/**
+ * The custom block bound to a workflow (with live-derived input fields), or `null`
+ * when the workflow isn't published as a block. One block per workflow is enforced
+ * at publish time. Used by the copilot deploy_custom_block tool.
+ */
+export async function getCustomBlockWithInputsByWorkflowId(
+  workflowId: string
+): Promise<CustomBlockWithInputs | null> {
+  const [row] = await db
+    .select({
+      block: customBlock,
+      workflowName: workflow.name,
+      workspaceId: workflow.workspaceId,
+      workspaceName: workspace.name,
+    })
+    .from(customBlock)
+    .innerJoin(workflow, eq(workflow.id, customBlock.workflowId))
+    .leftJoin(workspace, eq(workspace.id, workflow.workspaceId))
+    .where(eq(customBlock.workflowId, workflowId))
+    .limit(1)
+  return row ? hydrateCustomBlockRow(row) : null
+}
+
+/** Fetch a single custom block row by id. */
+export async function getCustomBlockById(id: string) {
+  const [row] = await db.select().from(customBlock).where(eq(customBlock.id, id)).limit(1)
+  return row ?? null
+}
+
+/**
+ * Org + source-workspace context for manage (edit/delete) authorization. Managing
+ * a block is gated on admin of its SOURCE workflow's workspace — the same workspace
+ * publishing required — so only an admin of the workspace that owns the workflow
+ * (or an org admin, who holds admin on every org workspace) can change its outputs.
+ * `null` when no block matches.
+ */
+export async function getCustomBlockManageContext(id: string): Promise<{
+  organizationId: string
+  sourceWorkspaceId: string | null
+  type: string
+  name: string
+} | null> {
+  const [row] = await db
+    .select({
+      organizationId: customBlock.organizationId,
+      sourceWorkspaceId: workflow.workspaceId,
+      type: customBlock.type,
+      name: customBlock.name,
+    })
+    .from(customBlock)
+    .innerJoin(workflow, eq(workflow.id, customBlock.workflowId))
+    .where(eq(customBlock.id, id))
+    .limit(1)
+  return row ?? null
+}
+
+/**
+ * Execution authority for a custom block, resolved by its block type. Used by the
+ * executor to run the bound workflow under the invocation-boundary model: the
+ * consumer needs no permission on the source workflow. Returns the authoritative
+ * `workflowId` from the DB (never trust a serialized value) plus the source
+ * workflow's **owner** (`workflow.userId`) — the same identity a normal deployed
+ * API/schedule/webhook run executes as. Using the owner (not the publisher) means
+ * the owner always has read on their own workflow, and owner deletion cascade-
+ * deletes the workflow → the custom_block row, so there is never an orphaned block.
+ * `null` when no enabled block matches the type.
+ */
+export async function getCustomBlockAuthority(
+  type: string,
+  consumerWorkspaceId: string | undefined
+): Promise<{
+  workflowId: string
+  organizationId: string
+  ownerUserId: string
+  exposedOutputs: CustomBlockOutput[]
+  /** Start-field ids (form keys) the publisher marked required. May reference removed fields. */
+  requiredInputIds: string[]
+} | null> {
+  // Scope resolution to the consumer's org: `(organizationId, type)` is the unique
+  // key, so without the org filter a `custom_block_*` type smuggled in from another
+  // org's serialized workflow could resolve and run that org's block.
+  if (!consumerWorkspaceId) return null
+  // Match `getCustomBlockRowsForWorkspace` (which builds the overlay) — include
+  // archived so a workspace that can serialize a custom block can also execute it,
+  // instead of failing mid-run with "no longer available".
+  const consumerWs = await getWorkspaceWithOwner(consumerWorkspaceId, { includeArchived: true })
+  if (!consumerWs?.organizationId) return null
+
+  const [row] = await db
+    .select({
+      workflowId: customBlock.workflowId,
+      organizationId: customBlock.organizationId,
+      enabled: customBlock.enabled,
+      outputs: customBlock.outputs,
+      inputs: customBlock.inputs,
+      ownerUserId: workflow.userId,
+    })
+    .from(customBlock)
+    .innerJoin(workflow, eq(workflow.id, customBlock.workflowId))
+    .where(
+      and(eq(customBlock.type, type), eq(customBlock.organizationId, consumerWs.organizationId))
+    )
+    .limit(1)
+
+  if (!row || !row.enabled) return null
+  return {
+    workflowId: row.workflowId,
+    organizationId: row.organizationId,
+    ownerUserId: row.ownerUserId,
+    exposedOutputs: row.outputs ?? [],
+    requiredInputIds: (row.inputs ?? []).filter((i) => i.required).map((i) => i.id),
+  }
+}
+
+export class CustomBlockValidationError extends Error {
+  constructor(message: string) {
+    super(message)
+    this.name = 'CustomBlockValidationError'
+  }
+}
+
+/**
+ * Publish a deployed workflow as an org-wide custom block. The source workflow
+ * must live in `workspaceId` — the workspace the caller was verified to admin —
+ * so a caller cannot publish another workspace's workflow (which then runs under
+ * that workspace owner's credentials and returns caller-chosen outputs). Also
+ * validates the workspace belongs to `organizationId` and the workflow is
+ * deployed, then inserts the row.
+ */
+export async function publishCustomBlock(params: {
+  organizationId: string
+  workspaceId: string
+  workflowId: string
+  userId: string
+  name: string
+  description: string
+  iconUrl?: string
+  inputs?: InputPlaceholder[]
+  exposedOutputs?: CustomBlockOutput[]
+}): Promise<CustomBlockWithInputs> {
+  const {
+    organizationId,
+    workspaceId,
+    workflowId,
+    userId,
+    name,
+    description,
+    iconUrl,
+    inputs,
+    exposedOutputs,
+  } = params
+
+  const [wf] = await db
+    .select({
+      id: workflow.id,
+      name: workflow.name,
+      workspaceId: workflow.workspaceId,
+      isDeployed: workflow.isDeployed,
+    })
+    .from(workflow)
+    .where(eq(workflow.id, workflowId))
+    .limit(1)
+
+  if (!wf) throw new CustomBlockValidationError('Workflow not found')
+  if (!wf.isDeployed) {
+    throw new CustomBlockValidationError('Workflow must be deployed before publishing as a block')
+  }
+
+  // Authorization boundary: the caller proved admin on `workspaceId` (route), so
+  // the source workflow must actually live there. Without this a workspace admin
+  // could publish a different workspace's workflow in the same org.
+  if (wf.workspaceId !== workspaceId) {
+    throw new CustomBlockValidationError('You can only publish a workflow from its own workspace')
+  }
+
+  const ws = wf.workspaceId ? await getWorkspaceWithOwner(wf.workspaceId) : null
+  if (!ws?.organizationId || ws.organizationId !== organizationId) {
+    throw new CustomBlockValidationError('Workflow does not belong to this organization')
+  }
+
+  // One block per workflow: the (org, type) unique index doesn't prevent the same
+  // workflow being published under a fresh `custom_block_*` type, so guard here.
+  const [existing] = await db
+    .select({ id: customBlock.id })
+    .from(customBlock)
+    .where(eq(customBlock.workflowId, workflowId))
+    .limit(1)
+  if (existing) {
+    throw new CustomBlockValidationError('This workflow is already published as a block')
+  }
+
+  const id = generateId()
+  const type = `${CUSTOM_BLOCK_TYPE_PREFIX}${generateShortId(10).toLowerCase()}`
+  const now = new Date()
+
+  await db.insert(customBlock).values({
+    id,
+    organizationId,
+    workflowId,
+    type,
+    name,
+    description,
+    iconUrl: iconUrl ?? null,
+    inputs: inputs ?? [],
+    outputs: exposedOutputs ?? [],
+    enabled: true,
+    createdBy: userId,
+    createdAt: now,
+    updatedAt: now,
+  })
+
+  logger.info('Published custom block', { id, type, organizationId, workflowId })
+
+  return {
+    id,
+    organizationId,
+    workflowId,
+    workflowName: wf.name,
+    workspaceId: wf.workspaceId,
+    workspaceName: ws?.name ?? null,
+    type,
+    name,
+    description,
+    iconUrl: iconUrl ?? null,
+    enabled: true,
+    inputFields: applyInputPlaceholders(inputs ?? null, await deriveInputFields(workflowId)),
+    exposedOutputs: exposedOutputs ?? [],
+  }
+}
+
+/**
+ * Update a custom block's presentation/enabled state. `iconUrl`: a URL
+ * sets/replaces the icon, `null` clears it (default icon), `undefined` leaves it
+ * unchanged.
+ */
+export async function updateCustomBlock(
+  id: string,
+  updates: {
+    name?: string
+    description?: string
+    enabled?: boolean
+    iconUrl?: string | null
+    inputs?: InputPlaceholder[]
+    exposedOutputs?: CustomBlockOutput[]
+  }
+): Promise<void> {
+  const patch: Partial<typeof customBlock.$inferInsert> = { updatedAt: new Date() }
+  if (updates.name !== undefined) patch.name = updates.name
+  if (updates.description !== undefined) patch.description = updates.description
+  if (updates.enabled !== undefined) patch.enabled = updates.enabled
+  if (updates.inputs !== undefined) patch.inputs = updates.inputs
+  if (updates.exposedOutputs !== undefined) patch.outputs = updates.exposedOutputs
+  if (updates.iconUrl !== undefined) patch.iconUrl = updates.iconUrl
+
+  await db.update(customBlock).set(patch).where(eq(customBlock.id, id))
+}
+
+/** Unpublish (hard-delete) a custom block. */
+export async function deleteCustomBlock(id: string): Promise<void> {
+  await db.delete(customBlock).where(eq(customBlock.id, id))
+}
+
+/**
+ * How many non-archived workflows in the org place the block, in their live
+ * editor state and/or their ACTIVE deployment snapshot. The two are scanned
+ * independently — a block removed in the editor can still ship in the active
+ * deployment (and vice versa), and the deployed placement is the one that
+ * actually runs. The deployment scan pre-filters with a raw-text match on the
+ * unique type slug so only near-exact matches pay the jsonb parse.
+ */
+export async function getCustomBlockUsageCounts(
+  organizationId: string,
+  blockType: string
+): Promise<{ usageCount: number; deployedUsageCount: number }> {
+  const orgActiveWorkflow = and(
+    eq(workspace.organizationId, organizationId),
+    isNull(workflow.archivedAt)
+  )
+  // Escape LIKE wildcards — the `_`s in `custom_block_<id>` would otherwise match
+  // any character and let unrelated states through to the jsonb parse.
+  const likePattern = `%${blockType.replace(/[\\%_]/g, '\\$&')}%`
+
+  const [liveRows, deployedRows] = await Promise.all([
+    db
+      .selectDistinct({ workflowId: workflow.id })
+      .from(workflowBlocks)
+      .innerJoin(workflow, eq(workflow.id, workflowBlocks.workflowId))
+      .innerJoin(workspace, eq(workspace.id, workflow.workspaceId))
+      .where(and(eq(workflowBlocks.type, blockType), orgActiveWorkflow)),
+    db
+      .select({ workflowId: workflow.id })
+      .from(workflowDeploymentVersion)
+      .innerJoin(workflow, eq(workflow.id, workflowDeploymentVersion.workflowId))
+      .innerJoin(workspace, eq(workspace.id, workflow.workspaceId))
+      .where(
+        and(
+          eq(workflowDeploymentVersion.isActive, true),
+          eq(workflow.isDeployed, true),
+          orgActiveWorkflow,
+          sql`${workflowDeploymentVersion.state}::text LIKE ${likePattern} ESCAPE '\\'`,
+          sql`EXISTS (
+            SELECT 1 FROM jsonb_each((${workflowDeploymentVersion.state})::jsonb -> 'blocks') AS b
+            WHERE b.value ->> 'type' = ${blockType}
+          )`
+        )
+      ),
+  ])
+
+  const usingWorkflowIds = new Set(liveRows.map((r) => r.workflowId))
+  for (const row of deployedRows) usingWorkflowIds.add(row.workflowId)
+
+  return { usageCount: usingWorkflowIds.size, deployedUsageCount: deployedRows.length }
+}