Merge branch 'feat/pull-a-variable' into feat/support-agent-sandbox

api/core/workflow/enums.py

@@ -64,6 +64,7 @@ class NodeType(StrEnum):
     TRIGGER_PLUGIN = "trigger-plugin"
     HUMAN_INPUT = "human-input"
     COMMAND = "command"
+    GROUP = "group"
 
     @property
     def is_trigger_node(self) -> bool:
@@ -255,6 +256,7 @@ class WorkflowNodeExecutionMetadataKey(StrEnum):
     LLM_CONTENT_SEQUENCE = "llm_content_sequence"
     LLM_TRACE = "llm_trace"
     COMPLETED_REASON = "completed_reason"  # completed reason for loop node
+    MENTION_PARENT_ID = "mention_parent_id"  # parent node id for extractor nodes
 
 
 class WorkflowNodeExecutionStatus(StrEnum):

api/core/workflow/graph/graph.py

@@ -307,7 +307,14 @@ class Graph:
         if not node_configs:
             raise ValueError("Graph must have at least one node")
 
-        node_configs = [node_config for node_config in node_configs if node_config.get("type", "") != "custom-note"]
+        # Filter out UI-only node types:
+        # - custom-note: top-level type (node_config.type == "custom-note")
+        # - group: data-level type (node_config.data.type == "group")
+        node_configs = [
+            node_config
+            for node_config in node_configs
+            if node_config.get("type", "") != "custom-note" and node_config.get("data", {}).get("type", "") != "group"
+        ]
 
         # Parse node configurations
         node_configs_map = cls._parse_node_configs(node_configs)

api/core/workflow/graph_engine/event_management/event_handlers.py

@@ -93,8 +93,8 @@ class EventHandler:
         Args:
             event: The event to handle
         """
-        # Events in loops or iterations are always collected
-        if event.in_loop_id or event.in_iteration_id:
+        # Events in loops, iterations, or extractor groups are always collected
+        if event.in_loop_id or event.in_iteration_id or event.in_mention_parent_id:
             self._event_collector.collect(event)
             return
         return self._dispatch(event)
@@ -125,6 +125,11 @@ class EventHandler:
         Args:
             event: The node started event
         """
+        # Check if this is an extractor node (has parent_node_id)
+        if self._is_extractor_node(event.node_id):
+            self._handle_extractor_node_started(event)
+            return
+
         # Track execution in domain model
         node_execution = self._graph_execution.get_or_create_node_execution(event.node_id)
         is_initial_attempt = node_execution.retry_count == 0
@@ -164,6 +169,11 @@ class EventHandler:
         Args:
             event: The node succeeded event
         """
+        # Check if this is an extractor node (has parent_node_id)
+        if self._is_extractor_node(event.node_id):
+            self._handle_extractor_node_success(event)
+            return
+
         # Update domain model
         node_execution = self._graph_execution.get_or_create_node_execution(event.node_id)
         node_execution.mark_taken()
@@ -226,6 +236,11 @@ class EventHandler:
         Args:
             event: The node failed event
         """
+        # Check if this is an extractor node (has parent_node_id)
+        if self._is_extractor_node(event.node_id):
+            self._handle_extractor_node_failed(event)
+            return
+
         # Update domain model
         node_execution = self._graph_execution.get_or_create_node_execution(event.node_id)
         node_execution.mark_failed(event.error)
@@ -345,3 +360,57 @@ class EventHandler:
                     self._graph_runtime_state.set_output("answer", value)
             else:
                 self._graph_runtime_state.set_output(key, value)
+
+    def _is_extractor_node(self, node_id: str) -> bool:
+        """
+        Check if node_id represents an extractor node (has parent_node_id).
+
+        Extractor nodes extract values from list[PromptMessage] for their parent node.
+        They have a parent_node_id field pointing to their parent node.
+        """
+        node = self._graph.nodes.get(node_id)
+        if node is None:
+            return False
+        return node.node_data.is_extractor_node
+
+    def _handle_extractor_node_started(self, event: NodeRunStartedEvent) -> None:
+        """
+        Handle extractor node started event.
+
+        Extractor nodes don't need full execution tracking, just collect the event.
+        """
+        # Track in response coordinator for stream ordering
+        self._response_coordinator.track_node_execution(event.node_id, event.id)
+
+        # Collect the event
+        self._event_collector.collect(event)
+
+    def _handle_extractor_node_success(self, event: NodeRunSucceededEvent) -> None:
+        """
+        Handle extractor node success event.
+
+        Extractor nodes need special handling:
+        - Store outputs in variable pool (for reference by other nodes)
+        - Accumulate token usage
+        - Collect the event for logging
+        - Do NOT process edges or enqueue next nodes (parent node handles that)
+        """
+        self._accumulate_node_usage(event.node_run_result.llm_usage)
+
+        # Store outputs in variable pool
+        self._store_node_outputs(event.node_id, event.node_run_result.outputs)
+
+        # Collect the event
+        self._event_collector.collect(event)
+
+    def _handle_extractor_node_failed(self, event: NodeRunFailedEvent) -> None:
+        """
+        Handle extractor node failed event.
+
+        Extractor node failures are collected for logging,
+        but the parent node is responsible for handling the error.
+        """
+        self._accumulate_node_usage(event.node_run_result.llm_usage)
+
+        # Collect the event for logging
+        self._event_collector.collect(event)

api/core/workflow/graph_engine/layers/persistence.py

@@ -68,6 +68,7 @@ class _NodeRuntimeSnapshot:
     predecessor_node_id: str | None
     iteration_id: str | None
     loop_id: str | None
+    mention_parent_id: str | None
     created_at: datetime
 
 
@@ -230,6 +231,7 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
         metadata = {
             WorkflowNodeExecutionMetadataKey.ITERATION_ID: event.in_iteration_id,
             WorkflowNodeExecutionMetadataKey.LOOP_ID: event.in_loop_id,
+            WorkflowNodeExecutionMetadataKey.MENTION_PARENT_ID: event.in_mention_parent_id,
         }
 
         domain_execution = WorkflowNodeExecution(
@@ -256,6 +258,7 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
             predecessor_node_id=event.predecessor_node_id,
             iteration_id=event.in_iteration_id,
             loop_id=event.in_loop_id,
+            mention_parent_id=event.in_mention_parent_id,
             created_at=event.start_at,
         )
         self._node_snapshots[event.id] = snapshot

api/core/workflow/graph_events/base.py

@@ -21,6 +21,12 @@ class GraphNodeEventBase(GraphEngineEvent):
     """iteration id if node is in iteration"""
     in_loop_id: str | None = None
     """loop id if node is in loop"""
+    in_mention_parent_id: str | None = None
+    """Parent node id if this is an extractor node event.
+    
+    When set, indicates this event belongs to an extractor node that
+    is extracting values for the specified parent node.
+    """
 
     # The version of the node, or "1" if not specified.
     node_version: str = "1"

api/core/workflow/nodes/agent/agent_node.py

@@ -12,11 +12,20 @@ from sqlalchemy.orm import Session
 from core.agent.entities import AgentToolEntity
 from core.agent.plugin_entities import AgentStrategyParameter
 from core.file import File, FileTransferMethod
+from core.memory.base import BaseMemory
+from core.memory.node_token_buffer_memory import NodeTokenBufferMemory
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance, ModelManager
 from core.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata
+from core.model_runtime.entities.message_entities import (
+    AssistantPromptMessage,
+    PromptMessage,
+    ToolPromptMessage,
+    UserPromptMessage,
+)
 from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from core.model_runtime.utils.encoders import jsonable_encoder
+from core.prompt.entities.advanced_prompt_entities import MemoryMode
 from core.provider_manager import ProviderManager
 from core.tools.entities.tool_entities import (
     ToolIdentity,
@@ -136,6 +145,9 @@ class AgentNode(Node[AgentNodeData]):
             )
             return
 
+        # Fetch memory for node memory saving
+        memory = self._fetch_memory_for_save()
+
         try:
             yield from self._transform_message(
                 messages=message_stream,
@@ -149,6 +161,7 @@ class AgentNode(Node[AgentNodeData]):
                 node_type=self.node_type,
                 node_id=self._node_id,
                 node_execution_id=self.id,
+                memory=memory,
             )
         except PluginDaemonClientSideError as e:
             transform_error = AgentMessageTransformError(
@@ -395,8 +408,20 @@ class AgentNode(Node[AgentNodeData]):
             icon = None
         return icon
 
-    def _fetch_memory(self, model_instance: ModelInstance) -> TokenBufferMemory | None:
-        # get conversation id
+    def _fetch_memory(self, model_instance: ModelInstance) -> BaseMemory | None:
+        """
+        Fetch memory based on configuration mode.
+
+        Returns TokenBufferMemory for conversation mode (default),
+        or NodeTokenBufferMemory for node mode (Chatflow only).
+        """
+        node_data = self.node_data
+        memory_config = node_data.memory
+
+        if not memory_config:
+            return None
+
+        # get conversation id (required for both modes in Chatflow)
         conversation_id_variable = self.graph_runtime_state.variable_pool.get(
             ["sys", SystemVariableKey.CONVERSATION_ID]
         )
@@ -404,16 +429,26 @@ class AgentNode(Node[AgentNodeData]):
             return None
         conversation_id = conversation_id_variable.value
 
-        with Session(db.engine, expire_on_commit=False) as session:
-            stmt = select(Conversation).where(Conversation.app_id == self.app_id, Conversation.id == conversation_id)
-            conversation = session.scalar(stmt)
-
-            if not conversation:
-                return None
-
-        memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)
-
-        return memory
+        # Return appropriate memory type based on mode
+        if memory_config.mode == MemoryMode.NODE:
+            # Node-level memory (Chatflow only)
+            return NodeTokenBufferMemory(
+                app_id=self.app_id,
+                conversation_id=conversation_id,
+                node_id=self._node_id,
+                tenant_id=self.tenant_id,
+                model_instance=model_instance,
+            )
+        else:
+            # Conversation-level memory (default)
+            with Session(db.engine, expire_on_commit=False) as session:
+                stmt = select(Conversation).where(
+                    Conversation.app_id == self.app_id, Conversation.id == conversation_id
+                )
+                conversation = session.scalar(stmt)
+                if not conversation:
+                    return None
+            return TokenBufferMemory(conversation=conversation, model_instance=model_instance)
 
     def _fetch_model(self, value: dict[str, Any]) -> tuple[ModelInstance, AIModelEntity | None]:
         provider_manager = ProviderManager()
@@ -457,6 +492,136 @@ class AgentNode(Node[AgentNodeData]):
         else:
             return [tool for tool in tools if tool.get("type") != ToolProviderType.MCP]
 
+    def _fetch_memory_for_save(self) -> BaseMemory | None:
+        """
+        Fetch memory instance for saving node memory.
+        This is a simplified version that doesn't require model_instance.
+        """
+        from core.model_manager import ModelManager
+        from core.model_runtime.entities.model_entities import ModelType
+
+        node_data = self.node_data
+        if not node_data.memory:
+            return None
+
+        # Get conversation_id
+        conversation_id_var = self.graph_runtime_state.variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID])
+        if not isinstance(conversation_id_var, StringSegment):
+            return None
+        conversation_id = conversation_id_var.value
+
+        # Return appropriate memory type based on mode
+        if node_data.memory.mode == MemoryMode.NODE:
+            # For node memory, we need a model_instance for token counting
+            # Use a simple default model for this purpose
+            try:
+                model_instance = ModelManager().get_default_model_instance(
+                    tenant_id=self.tenant_id,
+                    model_type=ModelType.LLM,
+                )
+            except Exception:
+                return None
+
+            return NodeTokenBufferMemory(
+                app_id=self.app_id,
+                conversation_id=conversation_id,
+                node_id=self._node_id,
+                tenant_id=self.tenant_id,
+                model_instance=model_instance,
+            )
+        else:
+            # Conversation-level memory doesn't need saving here
+            return None
+
+    def _build_context(
+        self,
+        parameters_for_log: dict[str, Any],
+        user_query: str,
+        assistant_response: str,
+        agent_logs: list[AgentLogEvent],
+    ) -> list[PromptMessage]:
+        """
+        Build context from user query, tool calls, and assistant response.
+        Format: user -> assistant(with tool_calls) -> tool -> assistant
+
+        The context includes:
+        - Current user query (always present, may be empty)
+        - Assistant message with tool_calls (if tools were called)
+        - Tool results
+        - Assistant's final response
+        """
+        context_messages: list[PromptMessage] = []
+
+        # Always add user query (even if empty, to maintain conversation structure)
+        context_messages.append(UserPromptMessage(content=user_query or ""))
+
+        # Extract actual tool calls from agent logs
+        # Only include logs with label starting with "CALL " - these are real tool invocations
+        tool_calls: list[AssistantPromptMessage.ToolCall] = []
+        tool_results: list[tuple[str, str, str]] = []  # (tool_call_id, tool_name, result)
+
+        for log in agent_logs:
+            if log.status == "success" and log.label and log.label.startswith("CALL "):
+                # Extract tool name from label (format: "CALL tool_name")
+                tool_name = log.label[5:]  # Remove "CALL " prefix
+                tool_call_id = log.message_id
+
+                # Parse tool response from data
+                data = log.data or {}
+                tool_response = ""
+
+                # Try to extract the actual tool response
+                if "tool_response" in data:
+                    tool_response = data["tool_response"]
+                elif "output" in data:
+                    tool_response = data["output"]
+                elif "result" in data:
+                    tool_response = data["result"]
+
+                if isinstance(tool_response, dict):
+                    tool_response = str(tool_response)
+
+                # Get tool input for arguments
+                tool_input = data.get("tool_call_input", {}) or data.get("input", {})
+                if isinstance(tool_input, dict):
+                    import json
+
+                    tool_input_str = json.dumps(tool_input, ensure_ascii=False)
+                else:
+                    tool_input_str = str(tool_input) if tool_input else ""
+
+                if tool_response:
+                    tool_calls.append(
+                        AssistantPromptMessage.ToolCall(
+                            id=tool_call_id,
+                            type="function",
+                            function=AssistantPromptMessage.ToolCall.ToolCallFunction(
+                                name=tool_name,
+                                arguments=tool_input_str,
+                            ),
+                        )
+                    )
+                    tool_results.append((tool_call_id, tool_name, str(tool_response)))
+
+        # Add assistant message with tool_calls if there were tool calls
+        if tool_calls:
+            context_messages.append(AssistantPromptMessage(content="", tool_calls=tool_calls))
+
+            # Add tool result messages
+            for tool_call_id, tool_name, result in tool_results:
+                context_messages.append(
+                    ToolPromptMessage(
+                        content=result,
+                        tool_call_id=tool_call_id,
+                        name=tool_name,
+                    )
+                )
+
+        # Add final assistant response
+        context_messages.append(AssistantPromptMessage(content=assistant_response))
+
+        return context_messages
+
     def _transform_message(
         self,
         messages: Generator[ToolInvokeMessage, None, None],
@@ -467,6 +632,7 @@ class AgentNode(Node[AgentNodeData]):
         node_type: NodeType,
         node_id: str,
         node_execution_id: str,
+        memory: BaseMemory | None = None,
     ) -> Generator[NodeEventBase, None, None]:
         """
         Convert ToolInvokeMessages into tuple[plain_text, files]
@@ -711,6 +877,12 @@ class AgentNode(Node[AgentNodeData]):
                 is_final=True,
             )
 
+        # Get user query from parameters for building context
+        user_query = parameters_for_log.get("query", "")
+
+        # Build context from history, user query, tool calls and assistant response
+        context = self._build_context(parameters_for_log, user_query, text, agent_logs)
+
         yield StreamCompletedEvent(
             node_run_result=NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
@@ -719,6 +891,7 @@ class AgentNode(Node[AgentNodeData]):
                     "usage": jsonable_encoder(llm_usage),
                     "files": ArrayFileSegment(value=files),
                     "json": json_output,
+                    "context": context,
                     **variables,
                 },
                 metadata={

api/core/workflow/nodes/base/__init__.py

@@ -1,4 +1,10 @@
-from .entities import BaseIterationNodeData, BaseIterationState, BaseLoopNodeData, BaseLoopState, BaseNodeData
+from .entities import (
+    BaseIterationNodeData,
+    BaseIterationState,
+    BaseLoopNodeData,
+    BaseLoopState,
+    BaseNodeData,
+)
 from .usage_tracking_mixin import LLMUsageTrackingMixin
 
 __all__ = [

api/core/workflow/nodes/base/entities.py

@@ -175,6 +175,16 @@ class BaseNodeData(ABC, BaseModel):
     default_value: list[DefaultValue] | None = None
     retry_config: RetryConfig = RetryConfig()
 
+    # Parent node ID when this node is used as an extractor.
+    # If set, this node is an "attached" extractor node that extracts values
+    # from list[PromptMessage] for the parent node's parameters.
+    parent_node_id: str | None = None
+
+    @property
+    def is_extractor_node(self) -> bool:
+        """Check if this node is an extractor node (has parent_node_id)."""
+        return self.parent_node_id is not None
+
     @property
     def default_value_dict(self) -> dict[str, Any]:
         if self.default_value:

api/core/workflow/nodes/base/node.py

@@ -273,10 +273,87 @@ class Node(Generic[NodeDataT]):
         """Check if execution should be stopped."""
         return self.graph_runtime_state.stop_event.is_set()
 
+    def _find_extractor_node_configs(self) -> list[dict[str, Any]]:
+        """
+        Find all extractor node configurations that have parent_node_id == self._node_id.
+
+        Returns:
+            List of node configuration dicts for extractor nodes
+        """
+        nodes = self.graph_config.get("nodes", [])
+        extractor_configs = []
+        for node_config in nodes:
+            node_data = node_config.get("data", {})
+            if node_data.get("parent_node_id") == self._node_id:
+                extractor_configs.append(node_config)
+        return extractor_configs
+
+    def _execute_mention_nodes(self) -> Generator[GraphNodeEventBase, None, None]:
+        """
+        Execute all extractor nodes associated with this node.
+
+        Extractor nodes are nodes with parent_node_id == self._node_id.
+        They are executed before the main node to extract values from list[PromptMessage].
+        """
+        from core.workflow.nodes.node_mapping import LATEST_VERSION, NODE_TYPE_CLASSES_MAPPING
+
+        extractor_configs = self._find_extractor_node_configs()
+        logger.debug("[Extractor] Found %d extractor nodes for parent '%s'", len(extractor_configs), self._node_id)
+        if not extractor_configs:
+            return
+
+        for config in extractor_configs:
+            node_id = config.get("id")
+            node_data = config.get("data", {})
+            node_type_str = node_data.get("type")
+
+            if not node_id or not node_type_str:
+                continue
+
+            # Get node class
+            try:
+                node_type = NodeType(node_type_str)
+            except ValueError:
+                continue
+
+            node_mapping = NODE_TYPE_CLASSES_MAPPING.get(node_type)
+            if not node_mapping:
+                continue
+
+            node_version = str(node_data.get("version", "1"))
+            node_cls = node_mapping.get(node_version) or node_mapping.get(LATEST_VERSION)
+            if not node_cls:
+                continue
+
+            # Instantiate and execute the extractor node
+            extractor_node = node_cls(
+                id=node_id,
+                config=config,
+                graph_init_params=self._graph_init_params,
+                graph_runtime_state=self.graph_runtime_state,
+            )
+
+            # Execute and process extractor node events
+            for event in extractor_node.run():
+                # Tag event with parent node id for stream ordering and history tracking
+                if isinstance(event, GraphNodeEventBase):
+                    event.in_mention_parent_id = self._node_id
+
+                if isinstance(event, NodeRunSucceededEvent):
+                    # Store extractor node outputs in variable pool
+                    outputs: Mapping[str, Any] = event.node_run_result.outputs
+                    for variable_name, variable_value in outputs.items():
+                        self.graph_runtime_state.variable_pool.add((node_id, variable_name), variable_value)
+                if not isinstance(event, NodeRunStreamChunkEvent):
+                    yield event
+
     def run(self) -> Generator[GraphNodeEventBase, None, None]:
         execution_id = self.ensure_execution_id()
         self._start_at = naive_utc_now()
 
+        # Step 1: Execute associated extractor nodes before main node execution
+        yield from self._execute_mention_nodes()
+
         # Create and push start event with required fields
         start_event = NodeRunStartedEvent(
             id=execution_id,

api/core/workflow/nodes/llm/entities.py

@@ -1,6 +1,6 @@
 import re
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal
+from typing import Annotated, Any, Literal, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
@@ -327,9 +327,28 @@ class ToolLogPayload(BaseModel):
         )
 
 
+class PromptMessageContext(BaseModel):
+    """Context variable reference in prompt template.
+
+    YAML/JSON format: { "$context": ["node_id", "variable_name"] }
+    This will be expanded to list[PromptMessage] at runtime.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    value_selector: Sequence[str] = Field(alias="$context")
+
+
+# Union type for prompt template items (static message or context variable reference)
+PromptTemplateItem: TypeAlias = Annotated[
+    LLMNodeChatModelMessage | PromptMessageContext,
+    Field(discriminator=None),
+]
+
+
 class LLMNodeData(BaseNodeData):
     model: ModelConfig
-    prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate
+    prompt_template: Sequence[PromptTemplateItem] | LLMNodeCompletionModelPromptTemplate
     prompt_config: PromptConfig = Field(default_factory=PromptConfig)
     memory: MemoryConfig | None = None
     context: ContextConfig

api/core/workflow/nodes/llm/llm_utils.py

@@ -8,12 +8,20 @@ from configs import dify_config
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.entities.provider_entities import ProviderQuotaType, QuotaUnit
 from core.file.models import File
-from core.memory.token_buffer_memory import TokenBufferMemory
+from core.memory import NodeTokenBufferMemory, TokenBufferMemory
+from core.memory.base import BaseMemory
 from core.model_manager import ModelInstance, ModelManager
 from core.model_runtime.entities.llm_entities import LLMUsage
+from core.model_runtime.entities.message_entities import (
+    AssistantPromptMessage,
+    MultiModalPromptMessageContent,
+    PromptMessage,
+    PromptMessageContentUnionTypes,
+    PromptMessageRole,
+)
 from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
-from core.prompt.entities.advanced_prompt_entities import MemoryConfig
+from core.prompt.entities.advanced_prompt_entities import MemoryConfig, MemoryMode
 from core.variables.segments import ArrayAnySegment, ArrayFileSegment, FileSegment, NoneSegment, StringSegment
 from core.workflow.enums import SystemVariableKey
 from core.workflow.nodes.llm.entities import ModelConfig
@@ -86,25 +94,56 @@ def fetch_files(variable_pool: VariablePool, selector: Sequence[str]) -> Sequenc
 
 
 def fetch_memory(
-    variable_pool: VariablePool, app_id: str, node_data_memory: MemoryConfig | None, model_instance: ModelInstance
-) -> TokenBufferMemory | None:
+    variable_pool: VariablePool,
+    app_id: str,
+    tenant_id: str,
+    node_data_memory: MemoryConfig | None,
+    model_instance: ModelInstance,
+    node_id: str = "",
+) -> BaseMemory | None:
+    """
+    Fetch memory based on configuration mode.
+
+    Returns TokenBufferMemory for conversation mode (default),
+    or NodeTokenBufferMemory for node mode (Chatflow only).
+
+    :param variable_pool: Variable pool containing system variables
+    :param app_id: Application ID
+    :param tenant_id: Tenant ID
+    :param node_data_memory: Memory configuration
+    :param model_instance: Model instance for token counting
+    :param node_id: Node ID in the workflow (required for node mode)
+    :return: Memory instance or None if not applicable
+    """
     if not node_data_memory:
         return None
 
-    # get conversation id
+    # Get conversation_id from variable pool (required for both modes in Chatflow)
     conversation_id_variable = variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID])
     if not isinstance(conversation_id_variable, StringSegment):
         return None
     conversation_id = conversation_id_variable.value
 
-    with Session(db.engine, expire_on_commit=False) as session:
-        stmt = select(Conversation).where(Conversation.app_id == app_id, Conversation.id == conversation_id)
-        conversation = session.scalar(stmt)
-        if not conversation:
+    # Return appropriate memory type based on mode
+    if node_data_memory.mode == MemoryMode.NODE:
+        # Node-level memory (Chatflow only)
+        if not node_id:
             return None
-
-    memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)
-    return memory
+        return NodeTokenBufferMemory(
+            app_id=app_id,
+            conversation_id=conversation_id,
+            node_id=node_id,
+            tenant_id=tenant_id,
+            model_instance=model_instance,
+        )
+    else:
+        # Conversation-level memory (default)
+        with Session(db.engine, expire_on_commit=False) as session:
+            stmt = select(Conversation).where(Conversation.app_id == app_id, Conversation.id == conversation_id)
+            conversation = session.scalar(stmt)
+            if not conversation:
+                return None
+        return TokenBufferMemory(conversation=conversation, model_instance=model_instance)
 
 
 def deduct_llm_quota(tenant_id: str, model_instance: ModelInstance, usage: LLMUsage):
@@ -170,3 +209,87 @@ def deduct_llm_quota(tenant_id: str, model_instance: ModelInstance, usage: LLMUs
                 )
                 session.execute(stmt)
                 session.commit()
+
+
+def build_context(
+    prompt_messages: Sequence[PromptMessage],
+    assistant_response: str,
+) -> list[PromptMessage]:
+    """
+    Build context from prompt messages and assistant response.
+    Excludes system messages and includes the current LLM response.
+    Returns list[PromptMessage] for use with ArrayPromptMessageSegment.
+
+    Note: Multi-modal content base64 data is truncated to avoid storing large data in context.
+    """
+    context_messages: list[PromptMessage] = [
+        _truncate_multimodal_content(m) for m in prompt_messages if m.role != PromptMessageRole.SYSTEM
+    ]
+    context_messages.append(AssistantPromptMessage(content=assistant_response))
+    return context_messages
+
+
+def _truncate_multimodal_content(message: PromptMessage) -> PromptMessage:
+    """
+    Truncate multi-modal content base64 data in a message to avoid storing large data.
+    Preserves the PromptMessage structure for ArrayPromptMessageSegment compatibility.
+
+    If file_ref is present, clears base64_data and url (they can be restored later).
+    Otherwise, truncates base64_data as fallback for legacy data.
+    """
+    content = message.content
+    if content is None or isinstance(content, str):
+        return message
+
+    # Process list content, handling multi-modal data based on file_ref availability
+    new_content: list[PromptMessageContentUnionTypes] = []
+    for item in content:
+        if isinstance(item, MultiModalPromptMessageContent):
+            if item.file_ref:
+                # Clear base64 and url, keep file_ref for later restoration
+                new_content.append(item.model_copy(update={"base64_data": "", "url": ""}))
+            else:
+                # Fallback: truncate base64_data if no file_ref (legacy data)
+                truncated_base64 = ""
+                if item.base64_data:
+                    truncated_base64 = item.base64_data[:10] + "...[TRUNCATED]..." + item.base64_data[-10:]
+                new_content.append(item.model_copy(update={"base64_data": truncated_base64}))
+        else:
+            new_content.append(item)
+
+    return message.model_copy(update={"content": new_content})
+
+
+def restore_multimodal_content_in_messages(messages: Sequence[PromptMessage]) -> list[PromptMessage]:
+    """
+    Restore multimodal content (base64 or url) in a list of PromptMessages.
+
+    When context is saved, base64_data is cleared to save storage space.
+    This function restores the content by parsing file_ref in each MultiModalPromptMessageContent.
+
+    Args:
+        messages: List of PromptMessages that may contain truncated multimodal content
+
+    Returns:
+        List of PromptMessages with restored multimodal content
+    """
+    from core.file import file_manager
+
+    return [_restore_message_content(msg, file_manager) for msg in messages]
+
+
+def _restore_message_content(message: PromptMessage, file_manager) -> PromptMessage:
+    """Restore multimodal content in a single PromptMessage."""
+    content = message.content
+    if content is None or isinstance(content, str):
+        return message
+
+    restored_content: list[PromptMessageContentUnionTypes] = []
+    for item in content:
+        if isinstance(item, MultiModalPromptMessageContent):
+            restored_item = file_manager.restore_multimodal_content(item)
+            restored_content.append(cast(PromptMessageContentUnionTypes, restored_item))
+        else:
+            restored_content.append(item)
+
+    return message.model_copy(update={"content": restored_content})

api/core/workflow/nodes/llm/node.py

@@ -7,7 +7,7 @@ import logging
 import re
 import time
 from collections.abc import Generator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from sqlalchemy import select
 
@@ -18,7 +18,7 @@ from core.file import File, FileTransferMethod, FileType, file_manager
 from core.helper.code_executor import CodeExecutor, CodeLanguage
 from core.llm_generator.output_parser.errors import OutputParserError
 from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output
-from core.memory.token_buffer_memory import TokenBufferMemory
+from core.memory.base import BaseMemory
 from core.model_manager import ModelInstance, ModelManager
 from core.model_runtime.entities import (
     ImagePromptMessageContent,
@@ -56,6 +56,7 @@ from core.tools.signature import sign_upload_file
 from core.tools.tool_manager import ToolManager
 from core.variables import (
     ArrayFileSegment,
+    ArrayPromptMessageSegment,
     ArraySegment,
     FileSegment,
     NoneSegment,
@@ -103,6 +104,7 @@ from .entities import (
     LLMTraceSegment,
     ModelConfig,
     ModelTraceSegment,
+    PromptMessageContext,
     StreamBuffers,
     ThinkTagStreamParser,
     ToolLogPayload,
@@ -181,8 +183,9 @@ class LLMNode(Node[LLMNodeData]):
         variable_pool = self.graph_runtime_state.variable_pool
 
         try:
-            # init messages template
-            self.node_data.prompt_template = self._transform_chat_messages(self.node_data.prompt_template)
+            # Parse prompt template to separate static messages and context references
+            prompt_template = self.node_data.prompt_template
+            static_messages, context_refs, template_order = self._parse_prompt_template()
 
             # fetch variables and fetch values from variable pool
             inputs = self._fetch_inputs(node_data=self.node_data)
@@ -230,8 +233,10 @@ class LLMNode(Node[LLMNodeData]):
             memory = llm_utils.fetch_memory(
                 variable_pool=variable_pool,
                 app_id=self.app_id,
+                tenant_id=self.tenant_id,
                 node_data_memory=self.node_data.memory,
                 model_instance=model_instance,
+                node_id=self._node_id,
             )
 
             query: str | None = None
@@ -242,21 +247,40 @@ class LLMNode(Node[LLMNodeData]):
                 ):
                     query = query_variable.text
 
-            prompt_messages, stop = LLMNode.fetch_prompt_messages(
-                sys_query=query,
-                sys_files=files,
-                context=context,
-                memory=memory,
-                model_config=model_config,
-                prompt_template=self.node_data.prompt_template,
-                memory_config=self.node_data.memory,
-                vision_enabled=self.node_data.vision.enabled,
-                vision_detail=self.node_data.vision.configs.detail,
-                variable_pool=variable_pool,
-                jinja2_variables=self.node_data.prompt_config.jinja2_variables,
-                tenant_id=self.tenant_id,
-                context_files=context_files,
-            )
+            # Get prompt messages
+            prompt_messages: Sequence[PromptMessage]
+            stop: Sequence[str] | None
+            if isinstance(prompt_template, list) and context_refs:
+                prompt_messages, stop = self._build_prompt_messages_with_context(
+                    context_refs=context_refs,
+                    template_order=template_order,
+                    static_messages=static_messages,
+                    query=query,
+                    files=files,
+                    context=context,
+                    memory=memory,
+                    model_config=model_config,
+                    context_files=context_files,
+                )
+            else:
+                prompt_messages, stop = LLMNode.fetch_prompt_messages(
+                    sys_query=query,
+                    sys_files=files,
+                    context=context,
+                    memory=memory,
+                    model_config=model_config,
+                    prompt_template=cast(
+                        Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
+                        self.node_data.prompt_template,
+                    ),
+                    memory_config=self.node_data.memory,
+                    vision_enabled=self.node_data.vision.enabled,
+                    vision_detail=self.node_data.vision.configs.detail,
+                    variable_pool=variable_pool,
+                    jinja2_variables=self.node_data.prompt_config.jinja2_variables,
+                    tenant_id=self.tenant_id,
+                    context_files=context_files,
+                )
 
             # Variables for outputs
             generation_data: LLMGenerationData | None = None
@@ -300,6 +324,7 @@ class LLMNode(Node[LLMNodeData]):
                     node_id=self._node_id,
                     node_type=self.node_type,
                     reasoning_format=self._node_data.reasoning_format,
+                    tenant_id=self.tenant_id,
                 )
 
             (
@@ -348,6 +373,7 @@ class LLMNode(Node[LLMNodeData]):
                 "reasoning_content": reasoning_content,
                 "usage": jsonable_encoder(usage),
                 "finish_reason": finish_reason,
+                "context": llm_utils.build_context(prompt_messages, clean_text),
             }
 
             # Build generation field
@@ -463,6 +489,7 @@ class LLMNode(Node[LLMNodeData]):
         node_id: str,
         node_type: NodeType,
         reasoning_format: Literal["separated", "tagged"] = "tagged",
+        tenant_id: str | None = None,
     ) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]:
         model_schema = model_instance.model_type_instance.get_model_schema(
             node_data_model.name, model_instance.credentials
@@ -486,6 +513,7 @@ class LLMNode(Node[LLMNodeData]):
                 stop=list(stop or []),
                 stream=True,
                 user=user_id,
+                tenant_id=tenant_id,
             )
         else:
             request_start_time = time.perf_counter()
@@ -743,6 +771,212 @@ class LLMNode(Node[LLMNodeData]):
 
         return messages
 
+    def _parse_prompt_template(
+        self,
+    ) -> tuple[list[LLMNodeChatModelMessage], list[PromptMessageContext], list[tuple[int, str]]]:
+        """
+        Parse prompt_template to separate static messages and context references.
+
+        Returns:
+            Tuple of (static_messages, context_refs, template_order)
+            - static_messages: list of LLMNodeChatModelMessage
+            - context_refs: list of PromptMessageContext
+            - template_order: list of (index, type) tuples preserving original order
+        """
+        prompt_template = self.node_data.prompt_template
+        static_messages: list[LLMNodeChatModelMessage] = []
+        context_refs: list[PromptMessageContext] = []
+        template_order: list[tuple[int, str]] = []
+
+        if isinstance(prompt_template, list):
+            for idx, item in enumerate(prompt_template):
+                if isinstance(item, PromptMessageContext):
+                    context_refs.append(item)
+                    template_order.append((idx, "context"))
+                else:
+                    static_messages.append(item)
+                    template_order.append((idx, "static"))
+            # Transform static messages for jinja2
+            if static_messages:
+                self.node_data.prompt_template = self._transform_chat_messages(static_messages)
+
+        return static_messages, context_refs, template_order
+
+    def _build_prompt_messages_with_context(
+        self,
+        *,
+        context_refs: list[PromptMessageContext],
+        template_order: list[tuple[int, str]],
+        static_messages: list[LLMNodeChatModelMessage],
+        query: str | None,
+        files: Sequence[File],
+        context: str | None,
+        memory: BaseMemory | None,
+        model_config: ModelConfigWithCredentialsEntity,
+        context_files: list[File],
+    ) -> tuple[list[PromptMessage], Sequence[str] | None]:
+        """
+        Build prompt messages by combining static messages and context references in DSL order.
+
+        Returns:
+            Tuple of (prompt_messages, stop_sequences)
+        """
+        variable_pool = self.graph_runtime_state.variable_pool
+
+        # Process messages in DSL order: iterate once and handle each type directly
+        combined_messages: list[PromptMessage] = []
+        context_idx = 0
+        static_idx = 0
+
+        for _, type_ in template_order:
+            if type_ == "context":
+                # Handle context reference
+                ctx_ref = context_refs[context_idx]
+                ctx_var = variable_pool.get(ctx_ref.value_selector)
+                if ctx_var is None:
+                    raise VariableNotFoundError(f"Variable {'.'.join(ctx_ref.value_selector)} not found")
+                if not isinstance(ctx_var, ArrayPromptMessageSegment):
+                    raise InvalidVariableTypeError(f"Variable {'.'.join(ctx_ref.value_selector)} is not array[message]")
+                # Restore multimodal content (base64/url) that was truncated when saving context
+                restored_messages = llm_utils.restore_multimodal_content_in_messages(ctx_var.value)
+                combined_messages.extend(restored_messages)
+                context_idx += 1
+            else:
+                # Handle static message
+                static_msg = static_messages[static_idx]
+                processed_msgs = LLMNode.handle_list_messages(
+                    messages=[static_msg],
+                    context=context,
+                    jinja2_variables=self.node_data.prompt_config.jinja2_variables or [],
+                    variable_pool=variable_pool,
+                    vision_detail_config=self.node_data.vision.configs.detail,
+                )
+                combined_messages.extend(processed_msgs)
+                static_idx += 1
+
+        # Append memory messages
+        memory_messages = _handle_memory_chat_mode(
+            memory=memory,
+            memory_config=self.node_data.memory,
+            model_config=model_config,
+        )
+        combined_messages.extend(memory_messages)
+
+        # Append current query if provided
+        if query:
+            query_message = LLMNodeChatModelMessage(
+                text=query,
+                role=PromptMessageRole.USER,
+                edition_type="basic",
+            )
+            query_msgs = LLMNode.handle_list_messages(
+                messages=[query_message],
+                context="",
+                jinja2_variables=[],
+                variable_pool=variable_pool,
+                vision_detail_config=self.node_data.vision.configs.detail,
+            )
+            combined_messages.extend(query_msgs)
+
+        # Handle files (sys_files and context_files)
+        combined_messages = self._append_files_to_messages(
+            messages=combined_messages,
+            sys_files=files,
+            context_files=context_files,
+            model_config=model_config,
+        )
+
+        # Filter empty messages and get stop sequences
+        combined_messages = self._filter_messages(combined_messages, model_config)
+        stop = self._get_stop_sequences(model_config)
+
+        return combined_messages, stop
+
+    def _append_files_to_messages(
+        self,
+        *,
+        messages: list[PromptMessage],
+        sys_files: Sequence[File],
+        context_files: list[File],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> list[PromptMessage]:
+        """Append sys_files and context_files to messages."""
+        vision_enabled = self.node_data.vision.enabled
+        vision_detail = self.node_data.vision.configs.detail
+
+        # Handle sys_files (will be deprecated later)
+        if vision_enabled and sys_files:
+            file_prompts = [
+                file_manager.to_prompt_message_content(file, image_detail_config=vision_detail) for file in sys_files
+            ]
+            if messages and isinstance(messages[-1], UserPromptMessage) and isinstance(messages[-1].content, list):
+                messages[-1] = UserPromptMessage(content=file_prompts + messages[-1].content)
+            else:
+                messages.append(UserPromptMessage(content=file_prompts))
+
+        # Handle context_files
+        if vision_enabled and context_files:
+            file_prompts = [
+                file_manager.to_prompt_message_content(file, image_detail_config=vision_detail)
+                for file in context_files
+            ]
+            if messages and isinstance(messages[-1], UserPromptMessage) and isinstance(messages[-1].content, list):
+                messages[-1] = UserPromptMessage(content=file_prompts + messages[-1].content)
+            else:
+                messages.append(UserPromptMessage(content=file_prompts))
+
+        return messages
+
+    def _filter_messages(
+        self, messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity
+    ) -> list[PromptMessage]:
+        """Filter empty messages and unsupported content types."""
+        filtered_messages: list[PromptMessage] = []
+
+        for message in messages:
+            if isinstance(message.content, list):
+                filtered_content: list[PromptMessageContentUnionTypes] = []
+                for content_item in message.content:
+                    # Skip non-text content if features are not defined
+                    if not model_config.model_schema.features:
+                        if content_item.type != PromptMessageContentType.TEXT:
+                            continue
+                        filtered_content.append(content_item)
+                        continue
+
+                    # Skip content if corresponding feature is not supported
+                    feature_map = {
+                        PromptMessageContentType.IMAGE: ModelFeature.VISION,
+                        PromptMessageContentType.DOCUMENT: ModelFeature.DOCUMENT,
+                        PromptMessageContentType.VIDEO: ModelFeature.VIDEO,
+                        PromptMessageContentType.AUDIO: ModelFeature.AUDIO,
+                    }
+                    required_feature = feature_map.get(content_item.type)
+                    if required_feature and required_feature not in model_config.model_schema.features:
+                        continue
+                    filtered_content.append(content_item)
+
+                # Simplify single text content
+                if len(filtered_content) == 1 and filtered_content[0].type == PromptMessageContentType.TEXT:
+                    message.content = filtered_content[0].data
+                else:
+                    message.content = filtered_content
+
+            if not message.is_empty():
+                filtered_messages.append(message)
+
+        if not filtered_messages:
+            raise NoPromptFoundError(
+                "No prompt found in the LLM configuration. "
+                "Please ensure a prompt is properly configured before proceeding."
+            )
+
+        return filtered_messages
+
+    def _get_stop_sequences(self, model_config: ModelConfigWithCredentialsEntity) -> Sequence[str] | None:
+        """Get stop sequences from model config."""
+        return model_config.stop
+
     def _fetch_jinja_inputs(self, node_data: LLMNodeData) -> dict[str, str]:
         variables: dict[str, Any] = {}
 
@@ -940,7 +1174,7 @@ class LLMNode(Node[LLMNodeData]):
         sys_query: str | None = None,
         sys_files: Sequence[File],
         context: str | None = None,
-        memory: TokenBufferMemory | None = None,
+        memory: BaseMemory | None = None,
         model_config: ModelConfigWithCredentialsEntity,
         prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
         memory_config: MemoryConfig | None = None,
@@ -2287,7 +2521,7 @@ def _calculate_rest_token(
 
 def _handle_memory_chat_mode(
     *,
-    memory: TokenBufferMemory | None,
+    memory: BaseMemory | None,
     memory_config: MemoryConfig | None,
     model_config: ModelConfigWithCredentialsEntity,
 ) -> Sequence[PromptMessage]:
@@ -2304,7 +2538,7 @@ def _handle_memory_chat_mode(
 
 def _handle_memory_completion_mode(
     *,
-    memory: TokenBufferMemory | None,
+    memory: BaseMemory | None,
     memory_config: MemoryConfig | None,
     model_config: ModelConfigWithCredentialsEntity,
 ) -> str:

api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py

@@ -7,7 +7,7 @@ from typing import Any, cast
 
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.file import File
-from core.memory.token_buffer_memory import TokenBufferMemory
+from core.memory.base import BaseMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities import ImagePromptMessageContent
 from core.model_runtime.entities.llm_entities import LLMUsage
@@ -145,8 +145,10 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         memory = llm_utils.fetch_memory(
             variable_pool=variable_pool,
             app_id=self.app_id,
+            tenant_id=self.tenant_id,
             node_data_memory=node_data.memory,
             model_instance=model_instance,
+            node_id=self._node_id,
         )
 
         if (
@@ -244,6 +246,10 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         # transform result into standard format
         result = self._transform_result(data=node_data, result=result or {})
 
+        # Build context from prompt messages and response
+        assistant_response = json.dumps(result, ensure_ascii=False)
+        context = llm_utils.build_context(prompt_messages, assistant_response)
+
         return NodeRunResult(
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
             inputs=inputs,
@@ -252,6 +258,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
                 "__is_success": 1 if not error else 0,
                 "__reason": error,
                 "__usage": jsonable_encoder(usage),
+                "context": context,
                 **result,
             },
             metadata={
@@ -299,7 +306,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         query: str,
         variable_pool: VariablePool,
         model_config: ModelConfigWithCredentialsEntity,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
     ) -> tuple[list[PromptMessage], list[PromptMessageTool]]:
@@ -381,7 +388,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         query: str,
         variable_pool: VariablePool,
         model_config: ModelConfigWithCredentialsEntity,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
@@ -419,7 +426,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         query: str,
         variable_pool: VariablePool,
         model_config: ModelConfigWithCredentialsEntity,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
@@ -453,7 +460,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         query: str,
         variable_pool: VariablePool,
         model_config: ModelConfigWithCredentialsEntity,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
@@ -681,7 +688,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         max_token_limit: int = 2000,
     ) -> list[ChatModelMessage]:
         model_mode = ModelMode(node_data.model.mode)
@@ -708,7 +715,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         max_token_limit: int = 2000,
     ):
         model_mode = ModelMode(node_data.model.mode)

api/core/workflow/nodes/question_classifier/question_classifier_node.py

@@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
-from core.memory.token_buffer_memory import TokenBufferMemory
+from core.memory.base import BaseMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities import LLMUsage, ModelPropertyKey, PromptMessageRole
 from core.model_runtime.utils.encoders import jsonable_encoder
@@ -96,8 +96,10 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         memory = llm_utils.fetch_memory(
             variable_pool=variable_pool,
             app_id=self.app_id,
+            tenant_id=self.tenant_id,
             node_data_memory=node_data.memory,
             model_instance=model_instance,
+            node_id=self._node_id,
         )
         # fetch instruction
         node_data.instruction = node_data.instruction or ""
@@ -197,10 +199,15 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
                 "model_provider": model_config.provider,
                 "model_name": model_config.model,
             }
+            # Build context from prompt messages and response
+            assistant_response = f"class_name: {category_name}, class_id: {category_id}"
+            context = llm_utils.build_context(prompt_messages, assistant_response)
+
             outputs = {
                 "class_name": category_name,
                 "class_id": category_id,
                 "usage": jsonable_encoder(usage),
+                "context": context,
             }
 
             return NodeRunResult(
@@ -312,7 +319,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         self,
         node_data: QuestionClassifierNodeData,
         query: str,
-        memory: TokenBufferMemory | None,
+        memory: BaseMemory | None,
         max_token_limit: int = 2000,
     ):
         model_mode = ModelMode(node_data.model.mode)

api/core/workflow/nodes/tool/entities.py

@@ -1,11 +1,63 @@
-from typing import Any, Literal, Union
+import re
+from collections.abc import Sequence
+from typing import Any, Literal, Self, Union
 
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, field_validator, model_validator
 from pydantic_core.core_schema import ValidationInfo
 
 from core.tools.entities.tool_entities import ToolProviderType
 from core.workflow.nodes.base.entities import BaseNodeData
 
+# Pattern to match mention value format: {{@node.context@}}instruction
+# The placeholder {{@node.context@}} must appear at the beginning
+# Format: {{@agent_node_id.context@}} where agent_node_id is dynamic, context is fixed
+MENTION_VALUE_PATTERN = re.compile(r"^\{\{@([a-zA-Z0-9_]+)\.context@\}\}(.*)$", re.DOTALL)
+
+
+def parse_mention_value(value: str) -> tuple[str, str]:
+    """Parse mention value into (node_id, instruction).
+
+    Args:
+        value: The mention value string like "{{@llm.context@}}extract keywords"
+
+    Returns:
+        Tuple of (node_id, instruction)
+
+    Raises:
+        ValueError: If value format is invalid
+    """
+    match = MENTION_VALUE_PATTERN.match(value)
+    if not match:
+        raise ValueError(
+            "For mention type, value must start with {{@node.context@}} placeholder, "
+            "e.g., '{{@llm.context@}}extract keywords'"
+        )
+    return match.group(1), match.group(2)
+
+
+class MentionConfig(BaseModel):
+    """Configuration for extracting value from context variable.
+
+    Used when a tool parameter needs to be extracted from list[PromptMessage]
+    context using an extractor LLM node.
+
+    Note: instruction is embedded in the value field as "{{@node.context@}}instruction"
+    """
+
+    # ID of the extractor LLM node
+    extractor_node_id: str
+
+    # Output variable selector from extractor node
+    # e.g., ["text"], ["structured_output", "query"]
+    output_selector: Sequence[str]
+
+    # Strategy when output is None
+    null_strategy: Literal["raise_error", "use_default"] = "raise_error"
+
+    # Default value when null_strategy is "use_default"
+    # Type should match the parameter's expected type
+    default_value: Any = None
+
 
 class ToolEntity(BaseModel):
     provider_id: str
@@ -35,7 +87,9 @@ class ToolNodeData(BaseNodeData, ToolEntity):
     class ToolInput(BaseModel):
         # TODO: check this type
         value: Union[Any, list[str]]
-        type: Literal["mixed", "variable", "constant"]
+        type: Literal["mixed", "variable", "constant", "mention"]
+        # Required config for mention type, extracting value from context variable
+        mention_config: MentionConfig | None = None
 
         @field_validator("type", mode="before")
         @classmethod
@@ -48,6 +102,9 @@ class ToolNodeData(BaseNodeData, ToolEntity):
 
             if typ == "mixed" and not isinstance(value, str):
                 raise ValueError("value must be a string")
+            elif typ == "mention":
+                # Skip here, will be validated in model_validator
+                pass
             elif typ == "variable":
                 if not isinstance(value, list):
                     raise ValueError("value must be a list")
@@ -58,6 +115,26 @@ class ToolNodeData(BaseNodeData, ToolEntity):
                 raise ValueError("value must be a string, int, float, bool or dict")
             return typ
 
+        @model_validator(mode="after")
+        def check_mention_type(self) -> Self:
+            """Validate mention type with mention_config."""
+            if self.type != "mention":
+                return self
+
+            value = self.value
+            if value is None:
+                return self
+
+            if not isinstance(value, str):
+                raise ValueError("value must be a string for mention type")
+            # For mention type, value must match format: {{@node.context@}}instruction
+            # This will raise ValueError if format is invalid
+            parse_mention_value(value)
+            # mention_config is required for mention type
+            if self.mention_config is None:
+                raise ValueError("mention_config is required for mention type")
+            return self
+
     tool_parameters: dict[str, ToolInput]
     # The version of the tool parameter.
     # If this value is None, it indicates this is a previous version

api/core/workflow/nodes/tool/tool_node.py

@@ -1,7 +1,10 @@
+import logging
 from collections.abc import Generator, Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import select
+
+logger = logging.getLogger(__name__)
 from sqlalchemy.orm import Session
 
 from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
@@ -184,6 +187,7 @@ class ToolNode(Node[ToolNodeData]):
             tool_parameters (Sequence[ToolParameter]): The list of tool parameters.
             variable_pool (VariablePool): The variable pool containing the variables.
             node_data (ToolNodeData): The data associated with the tool node.
+            for_log (bool): Whether to generate parameters for logging.
 
         Returns:
             Mapping[str, Any]: A dictionary containing the generated parameters.
@@ -199,14 +203,37 @@ class ToolNode(Node[ToolNodeData]):
                 continue
             tool_input = node_data.tool_parameters[parameter_name]
             if tool_input.type == "variable":
-                variable = variable_pool.get(tool_input.value)
+                if not isinstance(tool_input.value, list):
+                    raise ToolParameterError(f"Invalid variable selector for parameter '{parameter_name}'")
+                selector = tool_input.value
+                variable = variable_pool.get(selector)
                 if variable is None:
                     if parameter.required:
-                        raise ToolParameterError(f"Variable {tool_input.value} does not exist")
+                        raise ToolParameterError(f"Variable {selector} does not exist")
                     continue
                 parameter_value = variable.value
+            elif tool_input.type == "mention":
+                # Mention type: get value from extractor node's output
+                if tool_input.mention_config is None:
+                    raise ToolParameterError(
+                        f"mention_config is required for mention type parameter '{parameter_name}'"
+                    )
+                mention_config = tool_input.mention_config.model_dump()
+                try:
+                    parameter_value, found = variable_pool.resolve_mention(
+                        mention_config, parameter_name=parameter_name
+                    )
+                    if not found and parameter.required:
+                        raise ToolParameterError(
+                            f"Extractor output not found for required parameter '{parameter_name}'"
+                        )
+                    if not found:
+                        continue
+                except ValueError as e:
+                    raise ToolParameterError(str(e)) from e
             elif tool_input.type in {"mixed", "constant"}:
-                segment_group = variable_pool.convert_template(str(tool_input.value))
+                template = str(tool_input.value)
+                segment_group = variable_pool.convert_template(template)
                 parameter_value = segment_group.log if for_log else segment_group.text
             else:
                 raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'")
@@ -488,8 +515,12 @@ class ToolNode(Node[ToolNodeData]):
                 for selector in selectors:
                     result[selector.variable] = selector.value_selector
             elif input.type == "variable":
-                selector_key = ".".join(input.value)
-                result[f"#{selector_key}#"] = input.value
+                if isinstance(input.value, list):
+                    selector_key = ".".join(input.value)
+                    result[f"#{selector_key}#"] = input.value
+            elif input.type == "mention":
+                # Mention type: value is handled by extractor node, no direct variable reference
+                pass
             elif input.type == "constant":
                 pass
 

api/core/workflow/runtime/variable_pool.py

@@ -268,6 +268,58 @@ class VariablePool(BaseModel):
                 continue
             self.add(selector, value)
 
+    def resolve_mention(
+        self,
+        mention_config: Mapping[str, Any],
+        /,
+        *,
+        parameter_name: str = "",
+    ) -> tuple[Any, bool]:
+        """
+        Resolve a mention parameter value from an extractor node's output.
+
+        Mention parameters reference values extracted by an extractor LLM node
+        from list[PromptMessage] context.
+
+        Args:
+            mention_config: A dict containing:
+                - extractor_node_id: ID of the extractor LLM node
+                - output_selector: Selector path for the output variable (e.g., ["text"])
+                - null_strategy: "raise_error" or "use_default"
+                - default_value: Value to use when null_strategy is "use_default"
+            parameter_name: Name of the parameter being resolved (for error messages)
+
+        Returns:
+            Tuple of (resolved_value, found):
+                - resolved_value: The extracted value, or default_value if not found
+                - found: True if value was found, False if using default
+
+        Raises:
+            ValueError: If extractor_node_id is missing, or if null_strategy is
+                       "raise_error" and the value is not found
+        """
+        extractor_node_id = mention_config.get("extractor_node_id")
+        if not extractor_node_id:
+            raise ValueError(f"Missing extractor_node_id for mention parameter '{parameter_name}'")
+
+        output_selector = list(mention_config.get("output_selector", []))
+        null_strategy = mention_config.get("null_strategy", "raise_error")
+        default_value = mention_config.get("default_value")
+
+        # Build full selector: [extractor_node_id, ...output_selector]
+        full_selector = [extractor_node_id] + output_selector
+        variable = self.get(full_selector)
+
+        if variable is None:
+            if null_strategy == "use_default":
+                return default_value, False
+            raise ValueError(
+                f"Extractor node '{extractor_node_id}' output '{'.'.join(output_selector)}' "
+                f"not found for parameter '{parameter_name}'"
+            )
+
+        return variable.value, True
+
     @classmethod
     def empty(cls) -> VariablePool:
         """Create an empty variable pool."""