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

Code before:
_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()


def reset_context_provider() -> None:
    """Reset the capturer and all tenant-scoped context providers (primarily for tests)."""
    global _capturer
    _capturer = None
    _tenant_context_providers.clear()


Review comment (suggestion):
I suggest adding a dedicated exception type for this error.

Code after:

# 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.

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