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

Code before:
        """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}'")
    return prov()


def capture_current_context() -> IExecutionContext:
    """
    Capture current execution context from the calling environment.

    If a capturer is registered (e.g., Flask), use it. Otherwise, return a minimal
    context with NullAppContext + copy of current contextvars.
    """
    if _capturer is None:
        return ExecutionContext(
            app_context=NullAppContext(),
            context_vars=contextvars.copy_context(),
        )
    return _capturer()

Review comment (suggestion):
Preserve type safety for `read_context` and `register_context` is hard. I suggest adding a comment for the implementor of context values to provide type-safe wrappers for `read_context` and `register_context`. (Similar to the suggestion in Go's [context](https://pkg.go.dev/context#Context:~:text=//%20Packages%20that%20define%20a%20Context%20key%20should%20provide%20type%2Dsafe%20accessors%0A%09//%20for%20the%20values%20stored%20using%20that%20key%3A) pakcage).

Besides. I suggest recommending callers of `read_context` and `register_context` to  prefix their names to avoid name collision.

Code after:
        """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)."""
    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.

    Tip: use a namespaced "name" (e.g., "workflow.sandbox") to avoid key collisions.
    Consider adding a typed wrapper for this registration in your feature module.
    """
    _tenant_context_providers[(name, tenant_id)] = provider


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