jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(api):LLM node losing Flask context during parallel iterations (#26098)

Browse source
This commit is contained in:
quicksand 2025-09-24 10:09:35 +08:00 • committed by GitHub
commit a4acc64afd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 22 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

8
api/core/workflow/nodes/iteration/iteration_node.py
View file

@ -1,9 +1,11 @@
import contextvars
import logging import logging
from collections.abc import Generator, Mapping, Sequence from collections.abc import Generator, Mapping, Sequence
from concurrent.futures import Future, ThreadPoolExecutor, as_completed from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from datetime import UTC, datetime from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, NewType, cast from typing import TYPE_CHECKING, Any, NewType, cast
from flask import Flask, current_app
from typing_extensions import TypeIs from typing_extensions import TypeIs
from core.variables import IntegerVariable, NoneSegment from core.variables import IntegerVariable, NoneSegment
@ -35,6 +37,7 @@ from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
from core.workflow.nodes.base.node import Node from core.workflow.nodes.base.node import Node
from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
from libs.datetime_utils import naive_utc_now from libs.datetime_utils import naive_utc_now
from libs.flask_utils import preserve_flask_contexts
from .exc import ( from .exc import (
InvalidIteratorValueError, InvalidIteratorValueError,
@ -239,6 +242,8 @@ class IterationNode(Node):
self._execute_single_iteration_parallel, self._execute_single_iteration_parallel,
index=index, index=index,
item=item, item=item,
flask_app=current_app._get_current_object(), # type: ignore
context_vars=contextvars.copy_context(),
) )
future_to_index[future] = index future_to_index[future] = index
@ -281,8 +286,11 @@ class IterationNode(Node):
self, self,
index: int, index: int,
item: object, item: object,
flask_app: Flask,
context_vars: contextvars.Context,
) -> tuple[datetime, list[GraphNodeEventBase], object | None, int]: ) -> tuple[datetime, list[GraphNodeEventBase], object | None, int]:
"""Execute a single iteration in parallel mode and return results.""" """Execute a single iteration in parallel mode and return results."""
with preserve_flask_contexts(flask_app=flask_app, context_vars=context_vars):
iter_start_at = datetime.now(UTC).replace(tzinfo=None) iter_start_at = datetime.now(UTC).replace(tzinfo=None)
events: list[GraphNodeEventBase] = [] events: list[GraphNodeEventBase] = []
outputs_temp: list[object] = [] outputs_temp: list[object] = []