feat: expose current session to model plugins

src/dify_plugin/__init__.py

@@ -1,5 +1,6 @@
 from dify_plugin import _gevent  # noqa: F401 - import applies gevent patching
 from dify_plugin.config.config import DifyPluginEnv
+from dify_plugin.core.session_context import get_current_session
 from dify_plugin.interfaces.agent import AgentProvider, AgentStrategy
 from dify_plugin.interfaces.endpoint import Endpoint
 from dify_plugin.interfaces.model import ModelProvider
@@ -47,4 +48,5 @@
     "TextEmbeddingModel",
     "Tool",
     "ToolProvider",
+    "get_current_session",
 ]

src/dify_plugin/core/plugin_executor.py

@@ -53,6 +53,7 @@
 )
 from dify_plugin.core.plugin_registration import PluginRegistration
 from dify_plugin.core.runtime import Session
+from dify_plugin.core.session_context import use_current_session
 from dify_plugin.core.utils.http_parser import deserialize_request, serialize_response
 from dify_plugin.entities import ParameterOption
 from dify_plugin.entities.agent import AgentRuntime
@@ -242,22 +243,30 @@ def validate_model_credentials(
         return {"result": True, "credentials": data.credentials}
 
     def invoke_llm(self, session: Session, data: ModelInvokeLLMRequest) -> object:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, LargeLanguageModel):
-            return model_instance.invoke(
-                data.model,
-                data.credentials,
-                data.prompt_messages,
-                data.model_parameters,
-                data.tools,
-                data.stop,
-                data.stream,
-                data.user_id,
-            )
+            with use_current_session(session):
+                result = model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.prompt_messages,
+                    data.model_parameters,
+                    data.tools,
+                    data.stop,
+                    data.stream,
+                    data.user_id,
+                )
+            if not isinstance(result, Generator):
+                return result
+
+            def generator() -> Generator[object, None, None]:
+                with use_current_session(session):
+                    yield from result
+
+            return generator()
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -358,18 +367,18 @@ def invoke_text_embedding(
         session: Session,
         data: ModelInvokeTextEmbeddingRequest,
     ) -> object:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, TextEmbeddingModel):
-            return model_instance.invoke(
-                data.model,
-                data.credentials,
-                data.texts,
-                data.user_id,
-            )
+            with use_current_session(session):
+                return model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.texts,
+                    data.user_id,
+                )
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -380,19 +389,19 @@ def invoke_multimodal_embedding(
         session: Session,
         data: ModelInvokeMultimodalEmbeddingRequest,
     ) -> object:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, TextEmbeddingModel):
-            return model_instance.invoke_multimodal(
-                data.model,
-                data.credentials,
-                data.documents,
-                user=data.user_id,
-                input_type=data.input_type,
-            )
+            with use_current_session(session):
+                return model_instance.invoke_multimodal(
+                    data.model,
+                    data.credentials,
+                    data.documents,
+                    user=data.user_id,
+                    input_type=data.input_type,
+                )
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -422,21 +431,21 @@ def get_text_embedding_num_tokens(
         )
 
     def invoke_rerank(self, session: Session, data: ModelInvokeRerankRequest) -> object:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, RerankModel):
-            return model_instance.invoke(
-                data.model,
-                data.credentials,
-                data.query,
-                data.docs,
-                data.score_threshold,
-                data.top_n,
-                data.user_id,
-            )
+            with use_current_session(session):
+                return model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.query,
+                    data.docs,
+                    data.score_threshold,
+                    data.top_n,
+                    data.user_id,
+                )
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -447,21 +456,21 @@ def invoke_multimodal_rerank(
         session: Session,
         data: ModelInvokeMultimodalRerankRequest,
     ) -> object:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, RerankModel):
-            return model_instance.invoke_multimodal(
-                data.model,
-                data.credentials,
-                data.query,
-                data.docs,
-                score_threshold=data.score_threshold,
-                top_n=data.top_n,
-                user=data.user_id,
-            )
+            with use_current_session(session):
+                return model_instance.invoke_multimodal(
+                    data.model,
+                    data.credentials,
+                    data.query,
+                    data.docs,
+                    score_threshold=data.score_threshold,
+                    top_n=data.top_n,
+                    user=data.user_id,
+                )
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -472,26 +481,26 @@ def invoke_tts(
         session: Session,
         data: ModelInvokeTTSRequest,
     ) -> Generator[dict[str, str], None, None]:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
         if isinstance(model_instance, TTSModel):
-            b = model_instance.invoke(
-                data.model,
-                data.tenant_id,
-                data.credentials,
-                data.content_text,
-                data.voice,
-                data.user_id,
-            )
-            if isinstance(b, bytes | bytearray | memoryview):
-                yield {"result": binascii.hexlify(b).decode()}
-                return
+            with use_current_session(session):
+                b = model_instance.invoke(
+                    data.model,
+                    data.tenant_id,
+                    data.credentials,
+                    data.content_text,
+                    data.voice,
+                    data.user_id,
+                )
+                if isinstance(b, bytes | bytearray | memoryview):
+                    yield {"result": binascii.hexlify(b).decode()}
+                    return
 
-            for chunk in b:
-                yield {"result": binascii.hexlify(chunk).decode()}
+                for chunk in b:
+                    yield {"result": binascii.hexlify(chunk).decode()}
         else:
             msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
             raise TypeError(
@@ -526,7 +535,6 @@ def invoke_speech_to_text(
         session: Session,
         data: ModelInvokeSpeech2TextRequest,
     ) -> dict[str, str]:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
@@ -538,14 +546,14 @@ def invoke_speech_to_text(
 
             with pathlib.Path(temp.name).open("rb") as f:
                 if isinstance(model_instance, Speech2TextModel):
-                    return {
-                        "result": model_instance.invoke(
+                    with use_current_session(session):
+                        result = model_instance.invoke(
                             data.model,
                             data.credentials,
                             f,
                             data.user_id,
-                        ),
-                    }
+                        )
+                    return {"result": result}
                 msg = (
                     f"Model `{data.model_type}` not found for provider "
                     f"`{data.provider}`"
@@ -581,21 +589,20 @@ def invoke_moderation(
         session: Session,
         data: ModelInvokeModerationRequest,
     ) -> dict[str, bool]:
-        del session
         model_instance = self.registration.get_model_instance(
             data.provider,
             data.model_type,
         )
 
         if isinstance(model_instance, ModerationModel):
-            return {
-                "result": model_instance.invoke(
+            with use_current_session(session):
+                result = model_instance.invoke(
                     data.model,
                     data.credentials,
                     data.text,
                     data.user_id,
-                ),
-            }
+                )
+            return {"result": result}
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,

src/dify_plugin/core/session_context.py

@@ -0,0 +1,37 @@
+"""Current model invocation session access."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dify_plugin.core.runtime import Session
+
+_current_session: ContextVar[Session | None] = ContextVar(
+    "dify_plugin_current_session",
+    default=None,
+)
+
+
+def get_current_session() -> Session | None:
+    """Return the current model invocation session, if any.
+
+    Model plugins can read ``session.app_id`` from the returned session.
+    It is ``None`` for model calls outside an app execution context.
+
+    Returns:
+        The current session, or ``None``.
+    """
+    return _current_session.get()
+
+
+@contextmanager
+def use_current_session(session: Session) -> Iterator[None]:
+    token = _current_session.set(session)
+    try:
+        yield
+    finally:
+        _current_session.reset(token)

tests/core/__init__.py

@@ -0,0 +1 @@
+