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

[ryuta-kobayashi-ug]

Fixes #311

Summary

Make the request session (including app_id) accessible to model plugins during invocation, enabling provider-side cost attribution.

Parent issue: langgenius/dify#35772

Changes

• Add session_context.py with ContextVar-based get_current_session()
• In plugin_executor.py, set the ContextVar before model invocation and reset after Generator consumption
• Export get_current_session as a public API in __init__.py
• Add comprehensive tests (12 cases)

Design

Uses contextvars.ContextVar because:

• AIModel.__init__() is @final — cannot add session to constructor
• _invoke() is @abstractmethod — changing signature would break all existing plugins
• model_instance.invoke() returns a Generator — the ContextVar must remain set until the Generator is fully consumed (wrapped generator pattern)

Usage

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 request with app_id
            ...
        else:
            # out-of-app-context call, use 'dify_system' sentinel
            ...

Testing

• 12 unit tests covering ContextVar behavior, thread isolation, Generator lifecycle, and public API
• E2E tested with custom Bedrock plugin across all 5 app types

Compatibility

Fully backward compatible — existing model plugins require no changes.

[gemini-code-assist]

Code Review

This pull request implements request-scoped session context propagation for model plugins using Python's contextvars. It introduces a get_current_session utility and updates the PluginExecutor to manage session state across various model invocation types. A critical issue was identified in the invoke_llm implementation where managing the ContextVar outside of the generator's execution flow could lead to context leaks or ValueError exceptions. The reviewer suggests moving the context management logic entirely inside the generator wrapper to ensure proper cleanup.

[gemini-code-assist]

The current implementation of invoke_llm has a potential ContextVar leak and logic error.

1. Context Pollution: _current_session.set(session) is called at line 246, which is in the caller's context. Since reset() is only called when the generator is exhausted or closed, the ContextVar remains set in the current thread/greenlet even after invoke_llm returns, until that generator is fully consumed. This can lead to unexpected behavior if the caller performs other actions before iterating the generator.
2. Potential ValueError: If multiple generators are created in the same thread/greenlet before any are consumed, calling reset(token) in a different order than they were set will raise a ValueError because the token won't match the current state of the ContextVar stack.
3. Redundant try-except: The try...except at line 247 is mostly redundant because model_instance.invoke is a generator function; it returns a generator object immediately and won't raise exceptions until iteration starts.

To fix this, the set() and reset() calls should be moved entirely inside the generator's execution flow, similar to how invoke_tts is implemented.

            def _with_session_context() -> Generator:
                token = _current_session.set(session)
                try:
                    # invoke() returns a Generator (lazy). Keep ContextVar set
                    # until fully consumed — _invoke() needs session access.
                    yield from model_instance.invoke(
                        data.model,
                        data.credentials,
                        data.prompt_messages,
                        data.model_parameters,
                        data.tools,
                        data.stop,
                        data.stream,
                        data.user_id,
                    )
                finally:
                    _current_session.reset(token)

            return _with_session_context()

[ryuta-kobayashi-ug]

Thanks for the thorough review! You're right — the original implementation had context pollution and ordering issues.

Fixed in the latest push: moved set() and reset() entirely inside the generator, so the ContextVar is only active during iteration. This eliminates the leak window and the potential ValueError from out-of-order resets.