@@ -9,13 +9,21 @@ import {
99import type { ExecutionContext } from '@/executor/types'
1010import type { SerializedBlock } from '@/serializer/types'
1111
12- const { mockExecutorExecute, mockCreateSnapshot, mockResolveBillingAttribution, executorOptions } =
13- vi . hoisted ( ( ) => ( {
14- mockExecutorExecute : vi . fn ( ) ,
15- mockCreateSnapshot : vi . fn ( ) ,
16- mockResolveBillingAttribution : vi . fn ( ) ,
17- executorOptions : [ ] as Array < Record < string , any > > ,
18- } ) )
12+ const {
13+ mockExecutorExecute,
14+ mockCreateSnapshot,
15+ mockResolveBillingAttribution,
16+ mockGetCustomBlockAuthority,
17+ mockGetPersonalAndWorkspaceEnv,
18+ executorOptions,
19+ } = vi . hoisted ( ( ) => ( {
20+ mockExecutorExecute : vi . fn ( ) ,
21+ mockCreateSnapshot : vi . fn ( ) ,
22+ mockResolveBillingAttribution : vi . fn ( ) ,
23+ mockGetCustomBlockAuthority : vi . fn ( ) ,
24+ mockGetPersonalAndWorkspaceEnv : vi . fn ( ) ,
25+ executorOptions : [ ] as Array < Record < string , any > > ,
26+ } ) )
1927
2028vi . mock ( '@/executor' , ( ) => ( {
2129 Executor : class {
@@ -30,6 +38,14 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
3038 resolveBillingAttribution : mockResolveBillingAttribution ,
3139} ) )
3240
41+ vi . mock ( '@/lib/environment/utils' , ( ) => ( {
42+ getPersonalAndWorkspaceEnv : mockGetPersonalAndWorkspaceEnv ,
43+ } ) )
44+
45+ vi . mock ( '@/lib/workflows/custom-blocks/operations' , ( ) => ( {
46+ getCustomBlockAuthority : mockGetCustomBlockAuthority ,
47+ } ) )
48+
3349vi . mock ( '@/lib/logs/execution/snapshot/service' , ( ) => ( {
3450 snapshotService : { createSnapshotWithDeduplication : mockCreateSnapshot } ,
3551} ) )
@@ -321,6 +337,70 @@ describe('WorkflowBlockHandler', () => {
321337 expect ( mockResolveBillingAttribution ) . not . toHaveBeenCalled ( )
322338 } )
323339
340+ it ( 'resolves a source-scoped billing attribution for custom block children' , async ( ) => {
341+ const consumerAttribution = { actorUserId : 'consumer-1' , workspaceId : 'workspace-consumer' }
342+ const sourceAttribution = { actorUserId : 'owner-9' , workspaceId : 'workspace-source' }
343+ const customBlock = {
344+ ...mockBlock ,
345+ metadata : { id : 'custom_block_abc' , name : 'Published Block' } ,
346+ }
347+ const ctx = {
348+ ...mockContext ,
349+ workspaceId : 'workspace-consumer' ,
350+ metadata : { ...mockContext . metadata , billingAttribution : consumerAttribution } ,
351+ } as unknown as ExecutionContext
352+
353+ mockGetCustomBlockAuthority . mockResolvedValue ( {
354+ workflowId : 'source-workflow-id' ,
355+ organizationId : 'org-1' ,
356+ ownerUserId : 'owner-9' ,
357+ exposedOutputs : [ ] ,
358+ requiredInputIds : [ ] ,
359+ } )
360+ mockGetPersonalAndWorkspaceEnv . mockResolvedValue ( {
361+ personalDecrypted : { } ,
362+ workspaceDecrypted : { } ,
363+ } )
364+ mockResolveBillingAttribution . mockResolvedValue ( sourceAttribution )
365+ mockFetch . mockImplementation ( async ( url : unknown ) => {
366+ if ( String ( url ) . includes ( '/deployed' ) ) {
367+ return {
368+ ok : true ,
369+ json : ( ) =>
370+ Promise . resolve ( {
371+ data : {
372+ deployedState : { blocks : { } , edges : [ ] , loops : { } , parallels : { } } ,
373+ } ,
374+ } ) ,
375+ }
376+ }
377+ return {
378+ ok : true ,
379+ json : ( ) =>
380+ Promise . resolve ( {
381+ data : {
382+ name : 'Source Workflow' ,
383+ workspaceId : 'workspace-source' ,
384+ variables : { } ,
385+ } ,
386+ } ) ,
387+ }
388+ } )
389+ mockCreateSnapshot . mockResolvedValue ( { snapshot : { id : 'snapshot-1' } } )
390+ mockExecutorExecute . mockResolvedValue ( { success : true , output : { data : 'ok' } } )
391+
392+ await handler . execute ( ctx , customBlock , { } )
393+
394+ expect ( mockResolveBillingAttribution ) . toHaveBeenCalledWith ( {
395+ actorUserId : 'owner-9' ,
396+ workspaceId : 'workspace-source' ,
397+ } )
398+ expect ( executorOptions ) . toHaveLength ( 1 )
399+ expect ( executorOptions [ 0 ] . contextExtensions . billingAttribution ) . toBe ( sourceAttribution )
400+ expect ( executorOptions [ 0 ] . contextExtensions . userId ) . toBe ( 'owner-9' )
401+ expect ( executorOptions [ 0 ] . contextExtensions . workspaceId ) . toBe ( 'workspace-source' )
402+ } )
403+
324404 it ( 'should fail closed when the executing context has no workspace' , async ( ) => {
325405 mockFetch . mockResolvedValueOnce ( {
326406 ok : true ,
0 commit comments