ComfyUI_frontend/src/stores/subgraphStore.ts

import { defineStore } from 'pinia'
import { computed, ref } from 'vue'

import { t } from '@/i18n'
import { SubgraphNode } from '@/lib/litegraph/src/litegraph'
import { useSettingStore } from '@/platform/settings/settingStore'
import { useToastStore } from '@/platform/updates/common/toastStore'
import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
import type { LoadedComfyWorkflow } from '@/platform/workflow/management/stores/comfyWorkflow'
import { ComfyWorkflow } from '@/platform/workflow/management/stores/comfyWorkflow'
import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
import type {
  ComfyNode,
  ComfyWorkflowJSON,
  NodeId
} from '@/platform/workflow/validation/schemas/workflowSchema'
import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
import type { NodeError } from '@/schemas/apiSchema'
import type {
  ComfyNodeDef as ComfyNodeDefV1,
  InputSpec
} from '@/schemas/nodeDefSchema'
import { isCloud, isDesktop } from '@/platform/distribution/types'
import { TemplateIncludeOnDistributionEnum } from '@/platform/workflow/templates/types/template'
import { api } from '@/scripts/api'
import type { GlobalSubgraphData } from '@/scripts/api'
import { useDialogService } from '@/services/dialogService'
import { useExecutionStore } from '@/stores/executionStore'
import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
import type { UserFile } from '@/stores/userFileStore'

async function confirmOverwrite(name: string): Promise<boolean | null> {
  return await useDialogService().confirm({
    title: t('subgraphStore.overwriteBlueprintTitle'),
    type: 'overwriteBlueprint',
    message: t('subgraphStore.overwriteBlueprint'),
    itemList: [name]
  })
}

export const useSubgraphStore = defineStore('subgraph', () => {
  class SubgraphBlueprint extends ComfyWorkflow {
    static override readonly basePath = 'subgraphs/'
    override readonly tintCanvasBg = '#22227740'

    hasPromptedSave: boolean = false

    constructor(
      options: { path: string; modified: number; size: number },
      confirmFirstSave: boolean = false
    ) {
      super(options)
      this.hasPromptedSave = !confirmFirstSave
    }

    validateSubgraph() {
      if (!this.activeState?.definitions)
        throw new Error(
          'The root graph of a subgraph blueprint must consist of only a single subgraph node'
        )
      const { subgraphs } = this.activeState.definitions
      const { nodes } = this.activeState
      /**
       * Determines whether a given node represents a subgraph node.
       *
       * @param node - The node to test
       * @returns `true` if the node corresponds to a registered subgraph type, `false` otherwise.
       */
      function isSubgraphNode(node: ComfyNode) {
        return node && subgraphs.some((s: { id: string }) => s.id === node.type)
      }
      if (nodes.length == 1 && isSubgraphNode(nodes[0])) return
      const errors: Record<NodeId, NodeError> = {}
      //mark errors for all but first subgraph node
      let firstSubgraphFound = false
      for (let i = 0; i < nodes.length; i++) {
        if (!firstSubgraphFound && isSubgraphNode(nodes[i])) {
          firstSubgraphFound = true
          continue
        }
        errors[nodes[i].id] = {
          errors: [],
          class_type: nodes[i].type,
          dependent_outputs: []
        }
      }
      useExecutionStore().lastNodeErrors = errors
      useCanvasStore().getCanvas().draw(true, true)
      throw new Error(
        'The root graph of a subgraph blueprint must consist of only a single subgraph node'
      )
    }

    override async save(): Promise<UserFile> {
      this.validateSubgraph()
      if (
        !this.hasPromptedSave &&
        useSettingStore().get('Comfy.Workflow.WarnBlueprintOverwrite')
      ) {
        if (!(await confirmOverwrite(this.filename))) return this
        this.hasPromptedSave = true
      }
      // Extract metadata from subgraph.extra to workflow.extra before saving
      this.extractMetadataToWorkflowExtra()
      const ret = await super.save()
      // Force reload to update initialState with saved metadata
      registerNodeDef(await this.load({ force: true }), {
        category: 'Subgraph Blueprints/User'
      })
      return ret
    }

    /**
     * Moves all properties (except workflowRendererVersion) from subgraph.extra
     * to workflow.extra, then removes from subgraph.extra to avoid duplication.
     */
    private extractMetadataToWorkflowExtra(): void {
      if (!this.activeState) return
      const subgraph = this.activeState.definitions?.subgraphs?.[0]
      if (!subgraph?.extra) return

      const sgExtra = subgraph.extra as Record<string, unknown>
      const workflowExtra = (this.activeState.extra ??= {}) as Record<
        string,
        unknown
      >

      for (const key of Object.keys(sgExtra)) {
        if (key === 'workflowRendererVersion') continue
        workflowExtra[key] = sgExtra[key]
        delete sgExtra[key]
      }
    }

    override async saveAs(path: string) {
      this.validateSubgraph()
      this.hasPromptedSave = true
      // Extract metadata from subgraph.extra to workflow.extra before saving
      this.extractMetadataToWorkflowExtra()
      const ret = await super.saveAs(path)
      // Force reload to update initialState with saved metadata
      registerNodeDef(await this.load({ force: true }), {
        category: 'Subgraph Blueprints/User'
      })
      return ret
    }
    override async load({ force = false }: { force?: boolean } = {}): Promise<
      this & LoadedComfyWorkflow
    > {
      if (!force && this.isLoaded) return await super.load({ force })
      const loaded = await super.load({ force })
      const st = loaded.activeState
      const sg = (st.definitions?.subgraphs ?? []).find(
        (sg: { id: string }) => sg.id == st.nodes[0].type
      )
      if (!sg)
        throw new Error(
          'Loaded subgraph blueprint does not contain valid subgraph'
        )
      sg.name = st.nodes[0].title = this.filename

      // Copy blueprint metadata from workflow extra to subgraph extra
      // so it's available when editing via canvas.subgraph.extra
      if (st.extra) {
        const sgExtra = (sg.extra ??= {}) as Record<string, unknown>
        for (const [key, value] of Object.entries(st.extra)) {
          if (key === 'workflowRendererVersion') continue
          sgExtra[key] = value
        }
      }

      return loaded
    }
    override async promptSave(): Promise<string | null> {
      return await useDialogService().prompt({
        title: t('subgraphStore.saveBlueprint'),
        message: t('subgraphStore.blueprintNamePrompt'),
        defaultValue: this.filename
      })
    }
    override unload(): void {
      //Skip unloading. Even if a workflow is closed after editing,
      //it must remain loaded in order to be added to the graph
    }
  }
  const subgraphCache: Record<string, LoadedComfyWorkflow> = {}
  const typePrefix = 'SubgraphBlueprint.'
  const subgraphDefCache = ref<Map<string, ComfyNodeDefImpl>>(new Map())
  const canvasStore = useCanvasStore()
  const subgraphBlueprints = computed(() => [
    ...subgraphDefCache.value.values()
  ])
  /**
   * Loads user and installed subgraph blueprints, registers their node definitions, and attaches them to the workflow store.
   *
   * Loads blueprints from the user's SubgraphBlueprints folder and from the platform-provided (global) set, filters installed blueprints by distribution and custom-node requirements, registers resulting node definitions, and attaches the corresponding workflows to the workflow store. Any load failures are logged to the console and surfaced to the user as an error toast summarizing the number or list of failures.
   */
  async function fetchSubgraphs() {
    async function loadBlueprint(options: {
      path: string
      modified: number
      size: number
    }): Promise<void> {
      options.path = SubgraphBlueprint.basePath + options.path
      const bp = await new SubgraphBlueprint(options, true).load()
      useWorkflowStore().attachWorkflow(bp)
      registerNodeDef(bp, { category: 'Subgraph Blueprints/User' })
    }
    async function loadInstalledBlueprints() {
      async function loadGlobalBlueprint([k, v]: [string, GlobalSubgraphData]) {
        const path = SubgraphBlueprint.basePath + v.name + '.json'
        const blueprint = new SubgraphBlueprint({
          path,
          modified: Date.now(),
          size: -1
        })
        blueprint.originalContent = blueprint.content = await v.data
        blueprint.filename = v.name
        useWorkflowStore().attachWorkflow(blueprint)
        const loaded = await blueprint.load()
        const category = v.info.category
          ? `Subgraph Blueprints/${v.info.category}`
          : 'Subgraph Blueprints'
        registerNodeDef(
          loaded,
          {
            python_module: v.info.node_pack,
            display_name: v.name,
            category,
            search_aliases: v.info.search_aliases
          },
          k
        )
      }
      const subgraphs = await api.getGlobalSubgraphs()
      const currentDistribution: TemplateIncludeOnDistributionEnum = isCloud
        ? TemplateIncludeOnDistributionEnum.Cloud
        : isDesktop
          ? TemplateIncludeOnDistributionEnum.Desktop
          : TemplateIncludeOnDistributionEnum.Local
      const filteredEntries = Object.entries(subgraphs).filter(([, v]) => {
        if (!isCloud && (v.info.requiresCustomNodes?.length ?? 0) > 0)
          return false
        if (
          (v.info.includeOnDistributions?.length ?? 0) > 0 &&
          !v.info.includeOnDistributions!.includes(currentDistribution)
        )
          return false
        return true
      })
      await Promise.allSettled(filteredEntries.map(loadGlobalBlueprint))
    }

    const userSubs = (
      await api.listUserDataFullInfo(SubgraphBlueprint.basePath)
    ).filter((f) => f.path.endsWith('.json'))
    const settled = await Promise.allSettled([
      ...userSubs.map(loadBlueprint),
      loadInstalledBlueprints()
    ])

    const errors = settled
      .filter((i): i is PromiseRejectedResult => i.status === 'rejected')
      .map((i) => i.reason)
    errors.forEach((e) => console.error('Failed to load subgraph blueprint', e))
    if (errors.length > 0) {
      useToastStore().add({
        severity: 'error',
        summary: t('subgraphStore.loadFailure'),
        detail: errors.length > 3 ? `x${errors.length}` : `${errors}`,
        life: 6000
      })
    }
  }
  /**
   * Build and register a Comfy node definition for a loaded subgraph blueprint.
   *
   * Creates a ComfyNodeDefV1 from the blueprint's root subgraph node, applies any provided overrides,
   * stores the resulting node definition in the subgraph definition cache, and associates the
   * loaded workflow in the subgraph workflow cache under the given name.
   *
   * @param workflow - The loaded subgraph workflow to register as a node definition
   * @param overrides - Partial node definition fields to apply on top of the generated definition
   * @param name - The blueprint name to register (defaults to the workflow filename)
   * @throws Error if the workflow does not contain a root subgraph node suitable for registration
   */
  function registerNodeDef(
    workflow: LoadedComfyWorkflow,
    overrides: Partial<ComfyNodeDefV1> = {},
    name: string = workflow.filename
  ) {
    const subgraphNode = workflow.changeTracker.initialState.nodes[0]
    if (!subgraphNode) throw new Error('Invalid Subgraph Blueprint')
    subgraphNode.inputs ??= []
    subgraphNode.outputs ??= []
    //NOTE: Types are cast to string. This is only used for input coloring on previews
    const inputs = Object.fromEntries(
      subgraphNode.inputs.map(
        (i: { name: string; type: string | number | string[] }) => [
          i.name,
          [`${i.type}`, undefined] satisfies InputSpec
        ]
      )
    )
    const workflowExtra = workflow.initialState.extra
    const description =
      workflowExtra?.BlueprintDescription ?? 'User generated subgraph blueprint'
    const search_aliases = workflowExtra?.BlueprintSearchAliases
    const nodedefv1: ComfyNodeDefV1 = {
      input: { required: inputs },
      output: subgraphNode.outputs.map(
        (o: { type: string | number | string[] }) => `${o.type}`
      ),
      output_name: subgraphNode.outputs.map((o: { name: string }) => o.name),
      name: typePrefix + name,
      display_name: name,
      description,
      category: 'Subgraph Blueprints',
      output_node: false,
      python_module: 'blueprint',
      search_aliases,
      ...overrides
    }
    const nodeDefImpl = new ComfyNodeDefImpl(nodedefv1)
    subgraphDefCache.value.set(name, nodeDefImpl)
    subgraphCache[name] = workflow
  }
  async function publishSubgraph(providedName?: string) {
    const canvas = canvasStore.getCanvas()
    const subgraphNode = [...canvas.selectedItems][0]
    if (
      canvas.selectedItems.size !== 1 ||
      !(subgraphNode instanceof SubgraphNode)
    )
      throw new TypeError('Must have single SubgraphNode selected to publish')

    const { nodes = [], subgraphs = [] } = canvas._serializeItems([
      subgraphNode
    ])
    if (nodes.length != 1) {
      throw new TypeError('Must have single SubgraphNode selected to publish')
    }

    //create minimal workflow
    const workflowData = {
      revision: 0,
      last_node_id: subgraphNode.id,
      last_link_id: 0,
      nodes,
      links: [] as never[],
      version: 0.4,
      definitions: { subgraphs }
    }
    //prompt name
    const name =
      providedName ??
      (await useDialogService().prompt({
        title: t('subgraphStore.saveBlueprint'),
        message: t('subgraphStore.blueprintNamePrompt'),
        defaultValue: subgraphNode.title
      }))
    if (!name) return
    if (subgraphDefCache.value.has(name) && !(await confirmOverwrite(name)))
      //User has chosen not to overwrite.
      return

    //upload file
    const path = SubgraphBlueprint.basePath + name + '.json'
    const workflow = new SubgraphBlueprint({
      path,
      size: -1,
      modified: Date.now()
    })
    workflow.originalContent = JSON.stringify(workflowData)
    const loadedWorkflow = await workflow.load()
    //Mark non-temporary
    workflow.size = 1
    await workflow.save()
    //add to files list?
    useWorkflowStore().attachWorkflow(loadedWorkflow)
    useToastStore().add({
      severity: 'success',
      summary: t('subgraphStore.publishSuccess'),
      detail: t('subgraphStore.publishSuccessMessage'),
      life: 4000
    })
  }
  async function editBlueprint(nodeType: string) {
    const name = nodeType.slice(typePrefix.length)
    if (!(name in subgraphCache))
      //As loading is blocked on in startup, this can likely be changed to invalid type
      throw new Error('not yet loaded')
    useWorkflowStore().attachWorkflow(subgraphCache[name])
    await useWorkflowService().openWorkflow(subgraphCache[name])
    const canvas = useCanvasStore().getCanvas()
    if (canvas.graph && 'subgraph' in canvas.graph.nodes[0])
      canvas.setGraph(canvas.graph.nodes[0].subgraph)
  }
  function getBlueprint(nodeType: string): ComfyWorkflowJSON {
    const name = nodeType.slice(typePrefix.length)
    if (!(name in subgraphCache))
      //As loading is blocked on in startup, this can likely be changed to invalid type
      throw new Error('not yet loaded')
    return subgraphCache[name].changeTracker.initialState
  }
  async function deleteBlueprint(nodeType: string) {
    const name = nodeType.slice(typePrefix.length)
    if (!(name in subgraphCache)) throw new Error('not yet loaded')

    if (isGlobalBlueprint(name)) {
      useToastStore().add({
        severity: 'warn',
        summary: t('subgraphStore.cannotDeleteGlobal'),
        life: 4000
      })
      return
    }

    if (
      !(await useDialogService().confirm({
        title: t('subgraphStore.confirmDeleteTitle'),
        type: 'delete',
        message: t('subgraphStore.confirmDelete'),
        itemList: [name]
      }))
    )
      return

    await subgraphCache[name].delete()
    delete subgraphCache[name]
    subgraphDefCache.value.delete(name)
  }
  function isSubgraphBlueprint(
    workflow: unknown
  ): workflow is SubgraphBlueprint {
    return workflow instanceof SubgraphBlueprint
  }

  function isGlobalBlueprint(name: string): boolean {
    const nodeDef = subgraphDefCache.value.get(name)
    return nodeDef !== undefined && nodeDef.python_module !== 'blueprint'
  }

  return {
    deleteBlueprint,
    editBlueprint,
    fetchSubgraphs,
    getBlueprint,
    isGlobalBlueprint,
    isSubgraphBlueprint,
    publishSubgraph,
    subgraphBlueprints,
    typePrefix
  }
})