simstudioai--sim/apps/sim/executor/orchestrators/parallel.ts

import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { compactSubflowResults } from '@/lib/execution/payloads/serializer'
import { DEFAULTS } from '@/executor/constants'
import type { DAG } from '@/executor/dag/builder'
import type { EdgeManager } from '@/executor/execution/edge-manager'
import type { ParallelScope } from '@/executor/execution/state'
import type { BlockStateWriter, ContextExtensions } from '@/executor/execution/types'
import type { ExecutionContext, NormalizedBlockOutput } from '@/executor/types'
import { type ClonedSubflowInfo, ParallelExpander } from '@/executor/utils/parallel-expansion'
import {
  addSubflowErrorLog,
  buildBranchNodeId,
  buildParallelSentinelEndId,
  buildParallelSentinelStartId,
  buildSentinelEndId,
  buildSentinelStartId,
  emitSubflowSuccessEvents,
  extractBaseBlockId,
  extractBranchIndex,
} from '@/executor/utils/subflow-utils'
import { resolveArrayInputAsync } from '@/executor/utils/subflow-utils.server'
import type { VariableResolver } from '@/executor/variables/resolver'
import type { SerializedParallel } from '@/serializer/types'

const logger = createLogger('ParallelOrchestrator')
const DEFAULT_PARALLEL_BATCH_SIZE = 20

export interface ParallelBranchMetadata {
  branchIndex: number
  branchTotal: number
  distributionItem?: any
  parallelId: string
}

export interface ParallelAggregationResult {
  allBranchesComplete: boolean
  results?: unknown
  completedBranches?: number
  totalBranches?: number
}

export class ParallelOrchestrator {
  private expander = new ParallelExpander()

  constructor(
    private dag: DAG,
    private state: BlockStateWriter,
    private resolver: VariableResolver | null = null,
    private contextExtensions: ContextExtensions | null = null,
    private edgeManager: Pick<EdgeManager, 'clearDeactivatedEdgesForNodes'> | null = null
  ) {}

  async initializeParallelScope(ctx: ExecutionContext, parallelId: string): Promise<ParallelScope> {
    const parallelConfig = this.dag.parallelConfigs.get(parallelId)
    if (!parallelConfig) {
      throw new Error(`Parallel config not found: ${parallelId}`)
    }

    if (parallelConfig.nodes.length === 0) {
      const errorMessage =
        'Parallel has no executable blocks inside. Add or enable at least one block in the parallel.'
      logger.error(errorMessage, { parallelId })
      await this.addParallelErrorLog(ctx, parallelId, errorMessage, {})
      this.setErrorScope(ctx, parallelId, errorMessage)
      throw new Error(errorMessage)
    }

    let items: any[] | undefined
    let branchCount: number
    let isEmpty = false

    try {
      const resolved = await this.resolveBranchCount(ctx, parallelConfig, parallelId)
      branchCount = resolved.branchCount
      items = resolved.items
      isEmpty = resolved.isEmpty ?? false
    } catch (error) {
      const baseErrorMessage = toError(error).message
      const errorMessage = baseErrorMessage.startsWith('Parallel collection distribution is empty')
        ? baseErrorMessage
        : `Parallel Items did not resolve: ${baseErrorMessage}`
      logger.error(errorMessage, { parallelId, distribution: parallelConfig.distribution })
      await this.addParallelErrorLog(ctx, parallelId, errorMessage, {
        distribution: parallelConfig.distribution,
      })
      this.setErrorScope(ctx, parallelId, errorMessage)
      throw new Error(errorMessage)
    }

    if (isEmpty || branchCount === 0) {
      const scope: ParallelScope = {
        parallelId,
        totalBranches: 0,
        branchOutputs: new Map(),
        items: [],
        isEmpty: true,
      }

      if (!ctx.parallelExecutions) {
        ctx.parallelExecutions = new Map()
      }
      ctx.parallelExecutions.set(parallelId, scope)

      logger.info('Parallel scope initialized with empty distribution, skipping body', {
        parallelId,
        branchCount: 0,
      })

      return scope
    }

    const batchSize = this.resolveBatchSize(parallelConfig.batchSize)
    const currentBatchSize = Math.min(batchSize, branchCount)
    const scope: ParallelScope = {
      parallelId,
      totalBranches: branchCount,
      batchSize,
      currentBatchStart: 0,
      currentBatchSize,
      accumulatedOutputs: new Map(),
      branchOutputs: new Map(),
      items,
    }

    if (!ctx.parallelExecutions) {
      ctx.parallelExecutions = new Map()
    }
    ctx.parallelExecutions.set(parallelId, scope)

    logger.info('Parallel scope initialized', {
      parallelId,
      branchCount,
      batchSize,
      currentBatchSize,
    })

    return scope
  }

  prepareCurrentBatch(ctx: ExecutionContext, parallelId: string): void {
    const scope = ctx.parallelExecutions?.get(parallelId)
    if (!scope || scope.isEmpty) {
      return
    }

    const currentBatchStart = scope.currentBatchStart ?? 0
    const currentBatchSize =
      scope.currentBatchSize ??
      Math.min(scope.batchSize ?? DEFAULT_PARALLEL_BATCH_SIZE, scope.totalBranches)

    if (currentBatchSize <= 0) {
      return
    }

    const batchItems = scope.items?.slice(currentBatchStart, currentBatchStart + currentBatchSize)
    const { clonedSubflows, allBranchNodes, entryNodes, terminalNodes } =
      this.expander.expandParallel(this.dag, parallelId, currentBatchSize, batchItems, {
        branchIndexOffset: currentBatchStart,
        totalBranches: scope.totalBranches,
      })

    this.markRunFromBlockBatchDirty(
      ctx,
      parallelId,
      allBranchNodes,
      entryNodes,
      terminalNodes,
      clonedSubflows
    )
    this.registerClonedSubflows(ctx, parallelId, clonedSubflows)
    this.registerBranchMappings(ctx, parallelId, allBranchNodes)
    this.resetBatchExecutionState(allBranchNodes)
    this.edgeManager?.clearDeactivatedEdgesForNodes(new Set(allBranchNodes))

    logger.info('Prepared parallel batch', {
      parallelId,
      currentBatchStart,
      currentBatchSize,
      totalBranches: scope.totalBranches,
      branchNodeCount: allBranchNodes.length,
    })
  }

  private async resolveBranchCount(
    ctx: ExecutionContext,
    config: SerializedParallel,
    parallelId: string
  ): Promise<{ branchCount: number; items?: any[]; isEmpty?: boolean }> {
    if (config.parallelType === 'count') {
      return { branchCount: config.count ?? 1 }
    }

    const items = await this.resolveDistributionItems(ctx, config)
    if (items.length === 0) {
      logger.info('Parallel has empty distribution, skipping parallel body', { parallelId })
      return { branchCount: 0, items: [], isEmpty: true }
    }

    return { branchCount: items.length, items }
  }

  private async addParallelErrorLog(
    ctx: ExecutionContext,
    parallelId: string,
    errorMessage: string,
    inputData?: any
  ): Promise<void> {
    await addSubflowErrorLog(
      ctx,
      parallelId,
      'parallel',
      errorMessage,
      inputData || {},
      this.contextExtensions
    )
  }

  private setErrorScope(ctx: ExecutionContext, parallelId: string, errorMessage: string): void {
    const scope: ParallelScope = {
      parallelId,
      totalBranches: 0,
      branchOutputs: new Map(),
      items: [],
      validationError: errorMessage,
    }
    if (!ctx.parallelExecutions) {
      ctx.parallelExecutions = new Map()
    }
    ctx.parallelExecutions.set(parallelId, scope)
  }

  private async resolveDistributionItems(
    ctx: ExecutionContext,
    config: SerializedParallel
  ): Promise<any[]> {
    if (
      config.distribution === undefined ||
      config.distribution === null ||
      config.distribution === ''
    ) {
      throw new Error(
        'Parallel collection distribution is empty. Provide an array or a reference that resolves to a collection.'
      )
    }
    return resolveArrayInputAsync(
      ctx,
      config.distribution,
      this.resolver,
      buildParallelSentinelStartId(config.id)
    )
  }

  private resolveBatchSize(batchSize: unknown): number {
    const parsed =
      typeof batchSize === 'number' ? batchSize : Number.parseInt(String(batchSize), 10)
    if (Number.isNaN(parsed)) {
      return DEFAULT_PARALLEL_BATCH_SIZE
    }
    return Math.max(1, Math.min(DEFAULTS.MAX_PARALLEL_BRANCHES, parsed))
  }

  private registerClonedSubflows(
    ctx: ExecutionContext,
    parallelId: string,
    clonedSubflows: ClonedSubflowInfo[]
  ): void {
    if (clonedSubflows.length === 0 || !ctx.subflowParentMap) {
      return
    }

    const branchCloneMaps = new Map<number, Map<string, string>>()
    for (const clone of clonedSubflows) {
      let map = branchCloneMaps.get(clone.outerBranchIndex)
      if (!map) {
        map = new Map()
        branchCloneMaps.set(clone.outerBranchIndex, map)
      }
      map.set(clone.originalId, clone.clonedId)
    }

    for (const clone of clonedSubflows) {
      const originalEntry = ctx.subflowParentMap.get(clone.originalId)
      if (originalEntry) {
        const cloneMap = branchCloneMaps.get(clone.outerBranchIndex)
        const clonedParentId = cloneMap?.get(originalEntry.parentId)
        if (clonedParentId) {
          ctx.subflowParentMap.set(clone.clonedId, {
            parentId: clonedParentId,
            parentType: originalEntry.parentType,
            branchIndex: 0,
          })
        } else {
          ctx.subflowParentMap.set(clone.clonedId, {
            parentId: parallelId,
            parentType: 'parallel',
            branchIndex: clone.outerBranchIndex,
          })
        }
      } else {
        ctx.subflowParentMap.set(clone.clonedId, {
          parentId: parallelId,
          parentType: 'parallel',
          branchIndex: clone.outerBranchIndex,
        })
      }
    }
  }

  private markRunFromBlockBatchDirty(
    ctx: ExecutionContext,
    parallelId: string,
    allBranchNodes: string[],
    entryNodes: string[],
    terminalNodes: string[],
    clonedSubflows: ClonedSubflowInfo[]
  ): void {
    const dirtySet = ctx.runFromBlockContext?.dirtySet
    if (!dirtySet) return

    const parallelStartId = buildParallelSentinelStartId(parallelId)
    const parallelEndId = buildParallelSentinelEndId(parallelId)
    if (
      !dirtySet.has(parallelId) &&
      !dirtySet.has(parallelStartId) &&
      !dirtySet.has(parallelEndId)
    ) {
      return
    }

    for (const nodeId of allBranchNodes) dirtySet.add(nodeId)
    for (const nodeId of entryNodes) dirtySet.add(nodeId)
    for (const nodeId of terminalNodes) dirtySet.add(nodeId)

    const config = this.dag.parallelConfigs.get(parallelId)
    for (const nodeId of config?.nodes ?? []) {
      if (this.dag.parallelConfigs.has(nodeId) || this.dag.loopConfigs.has(nodeId)) {
        this.collectSubflowNodeIds(nodeId, dirtySet)
      }
    }

    for (const clone of clonedSubflows) {
      dirtySet.add(clone.clonedId)
      this.collectSubflowNodeIds(clone.clonedId, dirtySet)
    }
  }

  private collectSubflowNodeIds(
    subflowId: string,
    nodeIds: Set<string>,
    visited = new Set<string>()
  ): void {
    if (visited.has(subflowId)) return
    visited.add(subflowId)

    if (this.dag.parallelConfigs.has(subflowId)) {
      nodeIds.add(buildParallelSentinelStartId(subflowId))
      nodeIds.add(buildParallelSentinelEndId(subflowId))
      for (const childId of this.dag.parallelConfigs.get(subflowId)?.nodes ?? []) {
        if (this.dag.parallelConfigs.has(childId) || this.dag.loopConfigs.has(childId)) {
          this.collectSubflowNodeIds(childId, nodeIds, visited)
        } else {
          nodeIds.add(buildBranchNodeId(childId, 0))
        }
      }
      return
    }

    if (this.dag.loopConfigs.has(subflowId)) {
      nodeIds.add(buildSentinelStartId(subflowId))
      nodeIds.add(buildSentinelEndId(subflowId))
      for (const childId of this.dag.loopConfigs.get(subflowId)?.nodes ?? []) {
        if (this.dag.parallelConfigs.has(childId) || this.dag.loopConfigs.has(childId)) {
          this.collectSubflowNodeIds(childId, nodeIds, visited)
        } else {
          nodeIds.add(childId)
        }
      }
    }
  }

  /**
   * Stores a node's output in the branch outputs for later aggregation.
   * Aggregation is triggered by the sentinel-end node via the edge mechanism,
   * not by counting individual node completions. This avoids incorrect completion
   * detection when branches have conditional paths (error edges, conditions).
   */
  handleParallelBranchCompletion(
    ctx: ExecutionContext,
    parallelId: string,
    nodeId: string,
    output: NormalizedBlockOutput,
    branchIndexOverride?: number
  ): void {
    const scope = ctx.parallelExecutions?.get(parallelId)
    if (!scope) {
      logger.warn('Parallel scope not found for branch completion', { parallelId, nodeId })
      return
    }

    const mappedBranch = ctx.parallelBlockMapping?.get(nodeId)
    const branchIndex =
      branchIndexOverride ??
      (mappedBranch?.parallelId === parallelId
        ? mappedBranch.iterationIndex
        : (this.dag.nodes.get(nodeId)?.metadata.branchIndex ?? extractBranchIndex(nodeId)))
    if (branchIndex === null) {
      logger.warn('Could not extract branch index from node ID', { nodeId })
      return
    }

    if (!scope.branchOutputs.has(branchIndex)) {
      scope.branchOutputs.set(branchIndex, [])
    }
    scope.branchOutputs.get(branchIndex)!.push(output)
  }

  async aggregateParallelResults(
    ctx: ExecutionContext,
    parallelId: string
  ): Promise<ParallelAggregationResult> {
    const scope = ctx.parallelExecutions?.get(parallelId)
    if (!scope) {
      logger.error('Parallel scope not found for aggregation', { parallelId })
      return { allBranchesComplete: false }
    }

    const accumulatedOutputs =
      scope.accumulatedOutputs ?? new Map<number, NormalizedBlockOutput[]>()
    for (const [branchIndex, outputs] of scope.branchOutputs.entries()) {
      accumulatedOutputs.set(branchIndex, outputs)
    }
    scope.accumulatedOutputs = accumulatedOutputs
    scope.branchOutputs = new Map()

    const nextBatchStart =
      (scope.currentBatchStart ?? 0) + (scope.currentBatchSize ?? scope.totalBranches)
    if (nextBatchStart < scope.totalBranches) {
      /**
       * Compact accumulated outputs before scheduling the next batch. Each
       * block output is already individually compacted by `block-executor`, but
       * many below-threshold branch results can still exceed the aggregate
       * threshold over time. Re-running the existing subflow compactor over the
       * accumulated entries forces aggregate-size spills while existing
       * LargeValueRefs stay stable.
       */
      if (accumulatedOutputs.size > 0) {
        const accumulatedBranchIndexes = Array.from(accumulatedOutputs.keys()).sort((a, b) => a - b)
        const accumulatedResults = accumulatedBranchIndexes.map(
          (idx) => accumulatedOutputs.get(idx) ?? []
        )
        const compactedAccumulated = await compactSubflowResults(accumulatedResults, {
          workspaceId: ctx.workspaceId,
          workflowId: ctx.workflowId,
          executionId: ctx.executionId,
          largeValueExecutionIds: ctx.largeValueExecutionIds,
          largeValueKeys: ctx.largeValueKeys,
          allowLargeValueWorkflowScope: ctx.allowLargeValueWorkflowScope,
          userId: ctx.userId,
          requireDurable: true,
        })
        accumulatedBranchIndexes.forEach((branchIdx, position) => {
          accumulatedOutputs.set(branchIdx, compactedAccumulated[position])
        })
      }
      this.advanceToNextBatch(scope, nextBatchStart)
      return {
        allBranchesComplete: false,
        completedBranches: accumulatedOutputs.size,
        totalBranches: scope.totalBranches,
      }
    }

    const results: NormalizedBlockOutput[][] = []
    for (let i = 0; i < scope.totalBranches; i++) {
      const branchOutputs = accumulatedOutputs.get(i)
      if (!branchOutputs) {
        logger.warn('Missing branch output during parallel aggregation', { parallelId, branch: i })
      }
      results.push(branchOutputs ?? [])
    }
    const compactedResults = await compactSubflowResults(results, {
      workspaceId: ctx.workspaceId,
      workflowId: ctx.workflowId,
      executionId: ctx.executionId,
      largeValueExecutionIds: ctx.largeValueExecutionIds,
      largeValueKeys: ctx.largeValueKeys,
      allowLargeValueWorkflowScope: ctx.allowLargeValueWorkflowScope,
      userId: ctx.userId,
      requireDurable: true,
    })
    const output = { results: compactedResults }
    this.state.setBlockOutput(parallelId, output)
    scope.accumulatedOutputs = new Map()

    await emitSubflowSuccessEvents(ctx, parallelId, 'parallel', output, this.contextExtensions)

    return {
      allBranchesComplete: true,
      results: output.results,
      completedBranches: scope.totalBranches,
      totalBranches: scope.totalBranches,
    }
  }

  private advanceToNextBatch(scope: ParallelScope, nextBatchStart: number): void {
    const batchSize = scope.batchSize ?? DEFAULT_PARALLEL_BATCH_SIZE
    const remaining = scope.totalBranches - nextBatchStart
    const currentBatchSize = Math.min(batchSize, remaining)

    scope.currentBatchStart = nextBatchStart
    scope.currentBatchSize = currentBatchSize

    logger.info('Advanced to next parallel batch', {
      parallelId: scope.parallelId,
      nextBatchStart,
      currentBatchSize,
      totalBranches: scope.totalBranches,
    })
  }

  prepareForBatchContinuation(parallelId: string): void {
    this.state.unmarkExecuted(buildParallelSentinelStartId(parallelId))
    this.state.unmarkExecuted(buildParallelSentinelEndId(parallelId))
    this.state.deleteBlockState(buildParallelSentinelEndId(parallelId))
  }

  private resetBatchExecutionState(branchNodeIds: string[]): void {
    for (const nodeId of branchNodeIds) {
      const node = this.dag.nodes.get(nodeId)
      if (!node?.metadata.isParallelBranch) {
        continue
      }
      this.state.unmarkExecuted(nodeId)
      this.state.deleteBlockState(nodeId)
    }
  }

  private registerBranchMappings(
    ctx: ExecutionContext,
    parallelId: string,
    branchNodeIds: string[]
  ): void {
    if (branchNodeIds.length === 0) {
      return
    }

    if (!ctx.parallelBlockMapping) {
      ctx.parallelBlockMapping = new Map()
    }

    for (const nodeId of branchNodeIds) {
      const node = this.dag.nodes.get(nodeId)
      const branchIndex = node?.metadata.branchIndex ?? extractBranchIndex(nodeId)
      if (branchIndex === null || branchIndex === undefined) {
        continue
      }

      ctx.parallelBlockMapping.set(nodeId, {
        originalBlockId: node?.metadata.originalBlockId ?? extractBaseBlockId(nodeId),
        parallelId,
        iterationIndex: branchIndex,
      })
    }
  }

  extractBranchMetadata(nodeId: string): ParallelBranchMetadata | null {
    const node = this.dag.nodes.get(nodeId)
    if (!node?.metadata.isParallelBranch) {
      return null
    }

    const branchIndex = node.metadata.branchIndex ?? extractBranchIndex(nodeId)
    if (branchIndex === null) {
      return null
    }

    const parallelId =
      node.metadata.subflowType === 'parallel' ? node.metadata.subflowId : undefined
    if (!parallelId) {
      return null
    }

    return {
      branchIndex,
      branchTotal: node.metadata.branchTotal ?? 1,
      distributionItem: node.metadata.distributionItem,
      parallelId,
    }
  }

  getParallelScope(ctx: ExecutionContext, parallelId: string): ParallelScope | undefined {
    return ctx.parallelExecutions?.get(parallelId)
  }
}