Repository: langgenius/dify
File: api/core/workflow/context/execution_context.py (Python)

Code before:
    def with_app_context(self, app_context: AppContext) -> "ExecutionContextBuilder":
        """Set application context."""
        self._app_context = app_context
        return self

    def with_context_vars(self, context_vars: contextvars.Context) -> "ExecutionContextBuilder":
        """Set context variables."""
        self._context_vars = context_vars
        return self

    def with_user(self, user: Any) -> "ExecutionContextBuilder":
        """Set user."""
        self._user = user
        return self

    def build(self) -> ExecutionContext:
        """Build the execution context."""
        return ExecutionContext(
            app_context=self._app_context,
            context_vars=self._context_vars,
            user=self._user,
        )


_capturer: Callable[[], IExecutionContext] | None = None
_tenant_context_providers: dict[str, dict[str, Callable[[], BaseModel]]] = {}

T = TypeVar("T", bound=BaseModel)


def register_context_capturer(capturer: Callable[[], IExecutionContext]) -> None:
    """Register a single enterable execution context capturer (e.g., Flask)."""
    global _capturer
    _capturer = capturer


def register_context(name: str, tenant_id: str, provider: Callable[[], BaseModel]) -> None:
    """Register a tenant-specific provider for a named context."""
    _tenant_context_providers.setdefault(name, {})[tenant_id] = provider


def read_context(name: str, *, tenant_id: str) -> T:
    """
    Read a context value for a specific tenant.

    Raises KeyError if the provider for (name, tenant_id) is not registered.
    """
    providers = _tenant_context_providers.get(name, {})
    prov = providers.get(tenant_id)
    if prov is None:
        raise KeyError(f"Context provider '{name}' not registered for tenant '{tenant_id}'")

Review comment (suggestion):
It seems that is better to create `_tenant_context_providers` as a dictionary mapping from `(name, tenant_id)` tuple to the value.

The relationship between keys and values of the `_tenant_context_providers` should also be documented here.

Code after:
    def with_app_context(self, app_context: AppContext) -> "ExecutionContextBuilder":
        """Set application context."""
        self._app_context = app_context
        return self

    def with_context_vars(self, context_vars: contextvars.Context) -> "ExecutionContextBuilder":
        """Set context variables."""
        self._context_vars = context_vars
        return self

    def with_user(self, user: Any) -> "ExecutionContextBuilder":
        """Set user."""
        self._user = user
        return self

    def build(self) -> ExecutionContext:
        """Build the execution context."""
        return ExecutionContext(
            app_context=self._app_context,
            context_vars=self._context_vars,
            user=self._user,
        )


_capturer: Callable[[], IExecutionContext] | None = None

# Tenant-scoped providers using tuple keys for clarity and constant-time lookup.
# Key mapping:
#   (name, tenant_id) -> provider
# - name: namespaced identifier (recommend prefixing, e.g. "workflow.sandbox")
# - tenant_id: tenant identifier string
# Value:
#   provider: Callable[[], BaseModel] returning the typed context value
# Type-safety note:
#   - This registry cannot enforce that all providers for a given name return the same BaseModel type.
#   - Implementors SHOULD provide typed wrappers around register/read (like Go's context best practice),
#     e.g. def register_sandbox_ctx(tenant_id: str, p: Callable[[], SandboxContext]) and
#          def read_sandbox_ctx(tenant_id: str) -> SandboxContext.
_tenant_context_providers: dict[tuple[str, str], Callable[[], BaseModel]] = {}

T = TypeVar("T", bound=BaseModel)


class ContextProviderNotFoundError(KeyError):
    """Raised when a tenant-scoped context provider is missing for a given (name, tenant_id)."""

    pass


def register_context_capturer(capturer: Callable[[], IExecutionContext]) -> None:
    """Register a single enterable execution context capturer (e.g., Flask)."""