fix(custom-block): bill nested + failed-run hosted cost; expose real inputs to the agent

Author: TheodoreSpeaks

apps/sim/blocks/custom/server-overlay.ts

@@ -1,8 +1,12 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
+import type { WorkflowInputField } from '@/lib/workflows/input-format'
 import { buildCustomBlockConfig, type CustomBlockRow } from '@/blocks/custom/build-config'
 import { registerBlockOverlayResolver } from '@/blocks/custom/overlay'
 import type { BlockConfig, BlockIcon } from '@/blocks/types'
 
+/** A row for the overlay, optionally carrying live-derived Start input fields. */
+type CustomBlockOverlayRow = CustomBlockRow & { inputFields?: WorkflowInputField[] }
+
 /**
  * Server-side custom-block overlay. Resolves `custom_block_*` types during
  * serialization + execution from a per-request, per-org map held in
@@ -24,16 +28,23 @@ registerBlockOverlayResolver({
  * Run `fn` with the given org's custom blocks resolvable via `getBlock`/
  * `getAllBlocks`. Wrap every execution serializer entry point (execute route,
  * trigger.dev task, scheduled/webhook runs) at the org-context boundary so a
- * workflow containing a custom block can serialize and execute. Input mapping is
- * schema-agnostic, so the server needs no per-field editors (`inputFields: []`).
+ * workflow containing a custom block can serialize and execute.
+ *
+ * Execution passes bare rows: `inputMapping` is schema-agnostic, so no per-field
+ * editors are needed. Agent-facing callers (`get_blocks_metadata`, `edit_workflow`)
+ * pass rows carrying `inputFields` so `getBlock` exposes the real input sub-blocks —
+ * matching what the VFS block files show — instead of an empty schema.
  */
 export function withCustomBlockOverlay<T>(
-  rows: CustomBlockRow[],
+  rows: CustomBlockOverlayRow[],
   fn: () => Promise<T>
 ): Promise<T> {
   const map = new Map<string, BlockConfig>()
   for (const row of rows) {
-    map.set(row.type, buildCustomBlockConfig(row, [], { icon: PLACEHOLDER_ICON }))
+    map.set(
+      row.type,
+      buildCustomBlockConfig(row, row.inputFields ?? [], { icon: PLACEHOLDER_ICON })
+    )
   }
   return store.run(map, fn)
 }

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

@@ -3,6 +3,7 @@ import { getErrorMessage } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { getPersonalAndWorkspaceEnv } from '@/lib/environment/utils'
 import { buildNextCallChain, validateCallChain } from '@/lib/execution/call-chain'
+import { calculateCostSummary } from '@/lib/logs/execution/logging-factory'
 import { snapshotService } from '@/lib/logs/execution/snapshot/service'
 import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
 import type { TraceSpan } from '@/lib/logs/types'
@@ -58,6 +59,40 @@ function remapCustomBlockInputKeys(
   return remapped
 }
 
+/**
+ * Canonical hosted-key spend of a child run: the model/tool cost the way the
+ * parent bills it (recursing nested/iteration spans and de-duping model
+ * breakdowns), minus the base execution charge the parent applies once itself.
+ * A naive top-level `cost.total` sum undercounts when spend sits on nested children.
+ */
+function aggregateChildCost(childTraceSpans: TraceSpan[]): number {
+  if (childTraceSpans.length === 0) return 0
+  const summary = calculateCostSummary(childTraceSpans)
+  return Math.max(0, summary.totalCost - summary.baseExecutionCharge)
+}
+
+/**
+ * A single cost-only span so a FAILED custom block still bills the hosted-key spend
+ * its child already consumed (`block-executor` bills `error.childTraceSpans`),
+ * without exposing any of the source workflow's internal spans. Empty when free.
+ */
+function buildCostCarrierSpans(childCost: number, blockName: string, type: string): TraceSpan[] {
+  if (childCost <= 0) return []
+  const now = new Date().toISOString()
+  return [
+    {
+      id: generateId(),
+      name: blockName,
+      type,
+      duration: 0,
+      startTime: now,
+      endTime: now,
+      status: 'error',
+      cost: { total: childCost },
+    },
+  ]
+}
+
 type WorkflowTraceSpan = TraceSpan & {
   metadata?: Record<string, unknown>
   children?: WorkflowTraceSpan[]
@@ -328,7 +363,7 @@ export class WorkflowBlockHandler implements BlockHandler {
         // of the run's cost into billing — so roll their aggregate cost onto the
         // block itself. Custom blocks are org-scoped, so this bills the same org the
         // source workflow would bill if run directly, exactly as if it ran the key.
-        const childCost = childTraceSpans.reduce((sum, span) => sum + (span.cost?.total ?? 0), 0)
+        const childCost = aggregateChildCost(childTraceSpans)
         return this.projectCustomBlockOutput(executionResult, exposedOutputs, childCost)
       }
 
@@ -341,10 +376,29 @@ export class WorkflowBlockHandler implements BlockHandler {
       // blocks), trace spans, or execution result — the success path hides all of
       // these too. The real error is logged above for the publisher/ops; the
       // consumer gets only a generic failure attributed to the block they placed.
+      // But a child that failed AFTER consuming hosted keys still owes that spend,
+      // so capture the child's spans server-side, distill to the aggregate cost, and
+      // carry only that (no internals) so `block-executor` still bills it.
       if (isCustomBlock) {
+        let failedChildSpans: WorkflowTraceSpan[] = []
+        if (hasExecutionResult(error) && error.executionResult.logs) {
+          failedChildSpans = this.captureChildWorkflowLogs(
+            error.executionResult,
+            childWorkflowName,
+            ctx
+          )
+        } else if (ChildWorkflowError.isChildWorkflowError(error)) {
+          failedChildSpans = error.childTraceSpans
+        }
+        const blockName = block.metadata?.name || 'Custom block'
         throw new ChildWorkflowError({
           message: 'Custom block execution failed',
-          childWorkflowName: block.metadata?.name || 'Custom block',
+          childWorkflowName: blockName,
+          childTraceSpans: buildCostCarrierSpans(
+            aggregateChildCost(failedChildSpans),
+            blockName,
+            block.metadata?.id ?? 'custom_block'
+          ),
           childWorkflowInstanceId: instanceId,
         })
       }

apps/sim/lib/copilot/tools/server/router.ts

@@ -58,7 +58,7 @@ import { getCredentialsServerTool } from '@/lib/copilot/tools/server/user/get-cr
 import { setEnvironmentVariablesServerTool } from '@/lib/copilot/tools/server/user/set-environment-variables'
 import { editWorkflowServerTool } from '@/lib/copilot/tools/server/workflow/edit-workflow'
 import { queryLogsServerTool } from '@/lib/copilot/tools/server/workflow/query-logs'
-import { getCustomBlockRowsForWorkspace } from '@/lib/workflows/custom-blocks/operations'
+import { listCustomBlocksWithInputsForWorkspace } from '@/lib/workflows/custom-blocks/operations'
 import { withCustomBlockOverlay } from '@/blocks/custom/server-overlay'
 
 export type ExecuteResponseSuccess = z.output<typeof ExecuteResponseSuccessSchema>
@@ -248,7 +248,7 @@ export async function routeExecution(
   const result =
     CUSTOM_BLOCK_OVERLAY_TOOLS.has(toolName) && context?.workspaceId
       ? await withCustomBlockOverlay(
-          await getCustomBlockRowsForWorkspace(context.workspaceId),
+          await listCustomBlocksWithInputsForWorkspace(context.workspaceId),
           runTool
         )
       : await runTool()