feat: expose session context to model plugins via get_current_session()

Author: ryuta-kobayashi-ug

src/dify_plugin/__init__.py

@@ -4,6 +4,7 @@
 monkey.patch_all(sys=True)
 
 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
@@ -51,4 +52,5 @@
     "TextEmbeddingModel",
     "Tool",
     "ToolProvider",
+    "get_current_session",
 ]

src/dify_plugin/core/plugin_executor.py

@@ -51,6 +51,7 @@
 )
 from dify_plugin.core.plugin_registration import PluginRegistration
 from dify_plugin.core.runtime import Session
+from dify_plugin.core.session_context import _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,16 +243,31 @@ def invoke_llm(self, session: Session, data: ModelInvokeLLMRequest) -> object:
             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,
-            )
+            token = _current_session.set(session)
+            try:
+                result = model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.prompt_messages,
+                    data.model_parameters,
+                    data.tools,
+                    data.stop,
+                    data.stream,
+                    data.user_id,
+                )
+            except Exception:
+                _current_session.reset(token)
+                raise
+
+            # invoke() returns a Generator (lazy). Keep ContextVar set
+            # until fully consumed — _invoke() needs session access.
+            def _with_session_context(gen: Generator) -> Generator:
+                try:
+                    yield from gen
+                finally:
+                    _current_session.reset(token)
+
+            return _with_session_context(result)
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -291,12 +307,16 @@ def invoke_text_embedding(
             data.model_type,
         )
         if isinstance(model_instance, TextEmbeddingModel):
-            return model_instance.invoke(
-                data.model,
-                data.credentials,
-                data.texts,
-                data.user_id,
-            )
+            token = _current_session.set(session)
+            try:
+                return model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.texts,
+                    data.user_id,
+                )
+            finally:
+                _current_session.reset(token)
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -312,13 +332,17 @@ def invoke_multimodal_embedding(
             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,
-            )
+            token = _current_session.set(session)
+            try:
+                return model_instance.invoke_multimodal(
+                    data.model,
+                    data.credentials,
+                    data.documents,
+                    user=data.user_id,
+                    input_type=data.input_type,
+                )
+            finally:
+                _current_session.reset(token)
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -352,15 +376,19 @@ def invoke_rerank(self, session: Session, data: ModelInvokeRerankRequest) -> obj
             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,
-            )
+            token = _current_session.set(session)
+            try:
+                return model_instance.invoke(
+                    data.model,
+                    data.credentials,
+                    data.query,
+                    data.docs,
+                    data.score_threshold,
+                    data.top_n,
+                    data.user_id,
+                )
+            finally:
+                _current_session.reset(token)
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -376,15 +404,19 @@ def invoke_multimodal_rerank(
             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,
-            )
+            token = _current_session.set(session)
+            try:
+                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,
+                )
+            finally:
+                _current_session.reset(token)
         msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
         raise ValueError(
             msg,
@@ -400,20 +432,24 @@ def invoke_tts(
             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
+            token = _current_session.set(session)
+            try:
+                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()}
+            finally:
+                _current_session.reset(token)
         else:
             msg = f"Model `{data.model_type}` not found for provider `{data.provider}`"
             raise ValueError(
@@ -458,14 +494,18 @@ def invoke_speech_to_text(
 
             with pathlib.Path(temp.name).open("rb") as f:
                 if isinstance(model_instance, Speech2TextModel):
-                    return {
-                        "result": model_instance.invoke(
-                            data.model,
-                            data.credentials,
-                            f,
-                            data.user_id,
-                        ),
-                    }
+                    token = _current_session.set(session)
+                    try:
+                        return {
+                            "result": model_instance.invoke(
+                                data.model,
+                                data.credentials,
+                                f,
+                                data.user_id,
+                            ),
+                        }
+                    finally:
+                        _current_session.reset(token)
                 msg = (
                     f"Model `{data.model_type}` not found for provider "
                     f"`{data.provider}`"
@@ -506,14 +546,18 @@ def invoke_moderation(
         )
 
         if isinstance(model_instance, ModerationModel):
-            return {
-                "result": model_instance.invoke(
-                    data.model,
-                    data.credentials,
-                    data.text,
-                    data.user_id,
-                ),
-            }
+            token = _current_session.set(session)
+            try:
+                return {
+                    "result": model_instance.invoke(
+                        data.model,
+                        data.credentials,
+                        data.text,
+                        data.user_id,
+                    ),
+                }
+            finally:
+                _current_session.reset(token)
         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,59 @@
+"""
+Request-scoped session context for model plugins.
+
+Model plugins (LLM, Embedding, Rerank, etc.) do not receive the Session
+object through their ``_invoke()`` signature — unlike tool plugins which
+get it via their constructor.  This module bridges that gap by storing
+the current Session in a :class:`~contextvars.ContextVar` so that model
+plugin code can retrieve it on demand via :func:`get_current_session`.
+
+Usage in a custom model plugin::
+
+    from dify_plugin.core.session_context import get_current_session
+
+    class MyLLM(LargeLanguageModel):
+        def _invoke(self, model, credentials, prompt_messages, ...):
+            session = get_current_session()
+            if session and session.app_id:
+                # tag the request with the originating Dify app
+                ...
+
+Note on ``app_id`` being ``None``:
+
+    ``session.app_id`` is ``None`` when the model is invoked outside of
+    an app execution context — for example, RAG routing, conversation
+    title generation, or suggested question generation.  These calls
+    represent shared infrastructure costs not attributable to a specific
+    app.
+
+    When building provider-side cost dashboards, the recommended
+    approach is:
+
+    * If ``app_id`` is not ``None``, tag the request with it for
+      per-app cost attribution.
+    * If ``app_id`` is ``None``, either skip tagging or use a
+      sentinel value such as ``"dify_system"`` to bucket these
+      calls separately from external (non-Dify) traffic.
+"""
+
+from __future__ import annotations
+
+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(
+    "_current_session", default=None
+)
+
+
+def get_current_session() -> Session | None:
+    """Return the :class:`Session` for the current model invocation, or
+    ``None`` when called outside of a plugin dispatch context.
+
+    Returns:
+        The current session, or ``None``.
+    """
+    return _current_session.get()