stageflow-core 0.2.0__py3-none-any.whl

stageflow/__init__.py

@@ -0,0 +1,297 @@
+"""Stageflow - DAG-based pipeline orchestration framework.
+
+This package provides a framework for building observable, composable
+stage pipelines with parallel execution, cancellation, and interceptors.
+
+Core Components:
+- Stage: Protocol for pipeline stage implementations
+- Pipeline: Fluent builder for composing stages into DAGs
+- StageGraph: DAG executor with parallel execution
+- Interceptors: Middleware for cross-cutting concerns
+- EventSink: Protocol for event persistence
+
+Stage Kinds:
+- TRANSFORM: Data transformation stages (STT, TTS, LLM)
+- ENRICH: Context enrichment stages (Profile, Memory)
+- ROUTE: Routing decision stages (Router)
+- GUARD: Guardrail/validation stages
+- WORK: Side-effect stages (Persist, Assessment)
+- AGENT: Agentic/coaching stages
+
+Example:
+    from stageflow import Pipeline, Stage, StageOutput, StageKind
+
+    class MyStage:
+        name = "my_stage"
+        kind = StageKind.TRANSFORM
+
+        async def execute(self, ctx):
+            return StageOutput.ok(result="done")
+
+    pipeline = Pipeline().with_stage("my", MyStage, StageKind.TRANSFORM)
+    graph = pipeline.build()
+    results = await graph.run(ctx)
+
+Extension System:
+Stageflow provides a generic extension system for application-specific data.
+Use ContextSnapshot.extensions dict to store application data:
+
+    snapshot = ContextSnapshot(
+        ...
+        extensions={"skills": {"active_skill_ids": ["python"]}}
+    )
+
+For type-safe extensions, use the ExtensionRegistry in stageflow.extensions.
+"""
+
+# Core stage types
+# CLI and linting
+from stageflow.cli.lint import (
+    DependencyIssue,
+    DependencyLintResult,
+    IssueSeverity,
+    lint_pipeline,
+    lint_pipeline_file,
+)
+from stageflow.core import (
+    PipelineTimer,
+    Stage,
+    StageArtifact,
+    StageContext,
+    StageEvent,
+    StageKind,
+    StageOutput,
+    StageStatus,
+    create_stage_context,
+)
+
+# Events
+from stageflow.events import (
+    EventSink,
+    LoggingEventSink,
+    NoOpEventSink,
+    clear_event_sink,
+    get_event_sink,
+    set_event_sink,
+)
+
+# Extensions
+from stageflow.extensions import (
+    ExtensionHelper,
+    ExtensionRegistry,
+    TypedExtension,
+)
+
+# Observability protocols
+from stageflow.observability import (
+    CircuitBreaker,
+    CircuitBreakerOpenError,
+    PipelineRunLogger,
+    ProviderCallLogger,
+    error_summary_to_stages_patch,
+    error_summary_to_string,
+    get_circuit_breaker,
+    summarize_pipeline_error,
+)
+from stageflow.pipeline.dag import (
+    StageExecutionError,
+    StageGraph,
+    StageSpec,
+)
+
+# Interceptors
+from stageflow.pipeline.interceptors import (
+    BaseInterceptor,
+    ChildTrackerMetricsInterceptor,
+    CircuitBreakerInterceptor,
+    ErrorAction,
+    InterceptorContext,
+    InterceptorResult,
+    LoggingInterceptor,
+    MetricsInterceptor,
+    TimeoutInterceptor,
+    TracingInterceptor,
+    get_default_interceptors,
+    run_with_interceptors,
+)
+
+# Pipeline types
+from stageflow.pipeline.pipeline import (
+    Pipeline,
+    UnifiedStageSpec,
+)
+from stageflow.pipeline.registry import (
+    PipelineRegistry,
+    pipeline_registry,
+)
+from stageflow.pipeline.spec import (
+    CycleDetectedError,
+    PipelineValidationError,
+)
+from stageflow.pipeline.subpipeline import (
+    ChildRunTracker,
+    MaxDepthExceededError,
+    SubpipelineResult,
+    SubpipelineSpawner,
+    get_child_tracker,
+    get_subpipeline_spawner,
+)
+
+# Projector
+from stageflow.projector.service import (
+    WSMessageProjector,
+    WSMetadata,
+    WSOutboundMessage,
+    WSStatusUpdatePayload,
+    _coerce_uuid_str,
+)
+from stageflow.projector.service import (
+    WSMessageProjector as ProjectorService,  # Backward compatibility
+)
+
+# Protocols
+from stageflow.protocols import (
+    ConfigProvider,
+    CorrelationIds,
+    RunStore,
+)
+
+# Context types
+from stageflow.stages.context import (
+    PipelineContext,
+    extract_service,
+)
+from stageflow.stages.inputs import (
+    StageInputs,
+    create_stage_inputs,
+)
+from stageflow.stages.ports import (
+    AudioPorts,
+    CorePorts,
+    LLMPorts,
+    create_audio_ports,
+    create_core_ports,
+    create_llm_ports,
+)
+from stageflow.stages.result import (
+    StageError,
+    StageResult,
+)
+
+# Testing utilities (optional)
+from stageflow.testing import (
+    create_test_snapshot,
+    create_test_stage_context,
+)
+
+__all__ = [
+    # Core stage types
+    "Stage",
+    "StageKind",
+    "StageStatus",
+    "StageOutput",
+    "StageContext",
+    "StageArtifact",
+    "StageEvent",
+    "StageError",
+    "StageResult",
+    # Context utilities
+    "create_stage_context",
+    # Timer
+    "PipelineTimer",
+    # Pipeline types
+    "Pipeline",
+    "LinearPipeline",
+    "UnifiedStageSpec",
+    # DAG types
+    "StageExecutionError",
+    "StageGraph",
+    "StageSpec",
+    # Registry
+    "PipelineRegistry",
+    "pipeline_registry",
+    # Context types
+    "PipelineContext",
+    # Testing utilities
+    "create_test_snapshot",
+    "create_test_stage_context",
+    "StageError",
+    "extract_service",
+    # Interceptors
+    "BaseInterceptor",
+    "InterceptorResult",
+    "InterceptorContext",
+    "ErrorAction",
+    "LoggingInterceptor",
+    "MetricsInterceptor",
+    "ChildTrackerMetricsInterceptor",
+    "TracingInterceptor",
+    "CircuitBreakerInterceptor",
+    "TimeoutInterceptor",
+    "get_default_interceptors",
+    "run_with_interceptors",
+    # Events
+    "EventSink",
+    "NoOpEventSink",
+    "LoggingEventSink",
+    "get_event_sink",
+    "set_event_sink",
+    "clear_event_sink",
+    # Protocols
+    "RunStore",
+    "ConfigProvider",
+    "CorrelationIds",
+    # Observability
+    "CircuitBreaker",
+    "CircuitBreakerOpenError",
+    "PipelineRunLogger",
+    "ProviderCallLogger",
+    "summarize_pipeline_error",
+    "error_summary_to_string",
+    "error_summary_to_stages_patch",
+    "get_circuit_breaker",
+    # Extensions
+    "ExtensionRegistry",
+    "ExtensionHelper",
+    "TypedExtension",
+    # Stage inputs/ports
+    "StageInputs",
+    "create_stage_inputs",
+    "CorePorts",
+    "LLMPorts",
+    "AudioPorts",
+    "create_core_ports",
+    "create_llm_ports",
+    "create_audio_ports",
+    # Pipeline validation
+    "CycleDetectedError",
+    "PipelineValidationError",
+    # Subpipeline
+    "SubpipelineSpawner",
+    "SubpipelineResult",
+    "ChildRunTracker",
+    "MaxDepthExceededError",
+    "get_child_tracker",
+    "get_subpipeline_spawner",
+    # CLI and linting
+    "DependencyIssue",
+    "DependencyLintResult",
+    "IssueSeverity",
+    "lint_pipeline",
+    "lint_pipeline_file",
+    # Projector
+    "WSMessageProjector",
+    "ProjectorService",
+    "WSMetadata",
+    "WSOutboundMessage",
+    "WSStatusUpdatePayload",
+    "_coerce_uuid_str",
+    # Testing utilities
+    "create_test_snapshot",
+    "create_test_context",
+    "create_test_pipeline_context",
+    # Wide events
+    "WideEventEmitter",
+    "emit_stage_wide_event",
+    "emit_pipeline_wide_event",
+]

stageflow/agent/__init__.py

@@ -0,0 +1,19 @@
+"""Agent package containing ContextSnapshot and related types."""
+
+from stageflow.context import (
+    ContextSnapshot,
+    DocumentEnrichment,
+    MemoryEnrichment,
+    Message,
+    ProfileEnrichment,
+    RoutingDecision,
+)
+
+__all__ = [
+    "ContextSnapshot",
+    "Message",
+    "RoutingDecision",
+    "ProfileEnrichment",
+    "MemoryEnrichment",
+    "DocumentEnrichment",
+]

stageflow/auth/__init__.py

@@ -0,0 +1,56 @@
+"""Stageflow auth module - authentication and authorization types."""
+
+from stageflow.auth.context import AuthContext, OrgContext
+from stageflow.auth.errors import (
+    AuthenticationError,
+    CrossTenantAccessError,
+    InvalidTokenError,
+    MissingClaimsError,
+    TokenExpiredError,
+)
+from stageflow.auth.events import (
+    AuthFailureEvent,
+    AuthLoginEvent,
+    TenantAccessDeniedEvent,
+)
+from stageflow.auth.interceptors import (
+    AuthInterceptor,
+    JwtValidator,
+    MockJwtValidator,
+    OrgEnforcementInterceptor,
+)
+from stageflow.auth.tenant import (
+    TenantAwareLogger,
+    TenantContext,
+    TenantIsolationError,
+    TenantIsolationValidator,
+    clear_current_tenant,
+    get_current_tenant,
+    require_tenant,
+    set_current_tenant,
+)
+
+__all__ = [
+    "AuthContext",
+    "AuthenticationError",
+    "AuthFailureEvent",
+    "AuthInterceptor",
+    "AuthLoginEvent",
+    "CrossTenantAccessError",
+    "InvalidTokenError",
+    "JwtValidator",
+    "MissingClaimsError",
+    "MockJwtValidator",
+    "OrgContext",
+    "OrgEnforcementInterceptor",
+    "TenantAccessDeniedEvent",
+    "TenantAwareLogger",
+    "TenantContext",
+    "TenantIsolationError",
+    "TenantIsolationValidator",
+    "TokenExpiredError",
+    "clear_current_tenant",
+    "get_current_tenant",
+    "require_tenant",
+    "set_current_tenant",
+]

stageflow/auth/context.py

@@ -0,0 +1,119 @@
+"""Authentication and organization context types.
+
+This module provides immutable dataclasses for representing
+authenticated user context and organization context.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+from uuid import UUID
+
+
+@dataclass(frozen=True, slots=True)
+class AuthContext:
+    """Authenticated user context from JWT validation.
+
+    Immutable dataclass containing user identity and authorization
+    information extracted from a validated JWT token.
+
+    Attributes:
+        user_id: Unique user identifier
+        email: User's email address (optional)
+        org_id: Organization/tenant identifier (optional for personal accounts)
+        roles: Tuple of role names assigned to the user
+        session_id: Current session identifier
+    """
+
+    user_id: UUID
+    session_id: UUID
+    email: str | None = None
+    org_id: UUID | None = None
+    roles: tuple[str, ...] = field(default_factory=tuple)
+
+    def has_role(self, role: str) -> bool:
+        """Check if user has a specific role.
+
+        Args:
+            role: Role name to check
+
+        Returns:
+            True if user has the role
+        """
+        return role in self.roles
+
+    def is_admin(self) -> bool:
+        """Check if user has admin privileges.
+
+        Returns:
+            True if user has 'admin' or 'org_admin' role
+        """
+        return self.has_role("admin") or self.has_role("org_admin")
+
+    @property
+    def is_authenticated(self) -> bool:
+        """Check if this represents an authenticated user.
+
+        Returns:
+            Always True for a valid AuthContext
+        """
+        return True
+
+    def __repr__(self) -> str:
+        """Return string representation hiding sensitive email."""
+        email_display = f"{self.email[:3]}***" if self.email else None
+        return (
+            f"AuthContext(user_id={self.user_id!r}, "
+            f"email={email_display!r}, "
+            f"org_id={self.org_id!r}, "
+            f"roles={self.roles!r})"
+        )
+
+
+PlanTier = Literal["starter", "pro", "enterprise"]
+
+
+@dataclass(frozen=True, slots=True)
+class OrgContext:
+    """Organization context with plan and feature information.
+
+    Immutable dataclass containing organization metadata,
+    subscription tier, and enabled features.
+
+    Attributes:
+        org_id: Organization identifier
+        tenant_id: Tenant identifier (may differ from org_id in multi-tenant setups)
+        plan_tier: Subscription tier level
+        features: Tuple of enabled feature flags
+    """
+
+    org_id: UUID
+    tenant_id: UUID | None = None
+    plan_tier: PlanTier = "starter"
+    features: tuple[str, ...] = field(default_factory=tuple)
+
+    def has_feature(self, feature: str) -> bool:
+        """Check if organization has a specific feature enabled.
+
+        Args:
+            feature: Feature name to check
+
+        Returns:
+            True if feature is enabled
+        """
+        return feature in self.features
+
+    def __repr__(self) -> str:
+        return (
+            f"OrgContext(org_id={self.org_id!r}, "
+            f"plan_tier={self.plan_tier!r}, "
+            f"features={self.features!r})"
+        )
+
+
+__all__ = [
+    "AuthContext",
+    "OrgContext",
+    "PlanTier",
+]

stageflow/auth/errors.py

@@ -0,0 +1,87 @@
+"""Authentication and authorization exceptions.
+
+This module provides exception classes for authentication failures,
+token validation errors, and cross-tenant access violations.
+"""
+
+from __future__ import annotations
+
+
+class AuthenticationError(Exception):
+    """Base exception for authentication failures.
+
+    All authentication-related errors inherit from this class,
+    enabling catch-all handling for auth failures.
+    """
+
+    def __init__(self, message: str, code: str = "auth_error") -> None:
+        super().__init__(message)
+        self.code = code
+
+
+class TokenExpiredError(AuthenticationError):
+    """Raised when a JWT token has expired.
+
+    Attributes:
+        expired_at: ISO timestamp when the token expired
+    """
+
+    def __init__(self, message: str = "Token has expired", expired_at: str | None = None) -> None:
+        super().__init__(message, code="token_expired")
+        self.expired_at = expired_at
+
+
+class InvalidTokenError(AuthenticationError):
+    """Raised when a JWT token is invalid or malformed.
+
+    Attributes:
+        reason: Specific reason for invalidity (e.g., "invalid_signature", "malformed")
+    """
+
+    def __init__(self, message: str = "Invalid token", reason: str | None = None) -> None:
+        super().__init__(message, code="invalid_token")
+        self.reason = reason
+
+
+class MissingClaimsError(AuthenticationError):
+    """Raised when required JWT claims are missing.
+
+    Attributes:
+        missing_claims: List of claim names that are missing
+    """
+
+    def __init__(
+        self,
+        message: str = "Required claims are missing",
+        missing_claims: list[str] | None = None,
+    ) -> None:
+        super().__init__(message, code="missing_claims")
+        self.missing_claims = missing_claims or []
+
+
+class CrossTenantAccessError(AuthenticationError):
+    """Raised when a user attempts to access another tenant's resources.
+
+    Attributes:
+        user_org_id: The user's organization ID
+        resource_org_id: The resource's organization ID
+    """
+
+    def __init__(
+        self,
+        message: str = "Cross-tenant access denied",
+        user_org_id: str | None = None,
+        resource_org_id: str | None = None,
+    ) -> None:
+        super().__init__(message, code="cross_tenant_access")
+        self.user_org_id = user_org_id
+        self.resource_org_id = resource_org_id
+
+
+__all__ = [
+    "AuthenticationError",
+    "CrossTenantAccessError",
+    "InvalidTokenError",
+    "MissingClaimsError",
+    "TokenExpiredError",
+]

stageflow/auth/events.py

@@ -0,0 +1,118 @@
+"""Authentication audit event types.
+
+This module defines structured event types for authentication
+and authorization audit logging.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from typing import Any
+from uuid import UUID
+
+
+@dataclass(frozen=True, slots=True)
+class AuthLoginEvent:
+    """Event emitted on successful authentication.
+
+    Attributes:
+        user_id: Authenticated user ID
+        org_id: Organization ID (if applicable)
+        session_id: New session ID
+        timestamp: When authentication occurred
+        request_id: Request correlation ID
+        pipeline_run_id: Pipeline run correlation ID
+    """
+
+    user_id: UUID
+    session_id: UUID
+    org_id: UUID | None = None
+    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
+    request_id: UUID | None = None
+    pipeline_run_id: UUID | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for event emission."""
+        return {
+            "type": "auth.login",
+            "user_id": str(self.user_id),
+            "org_id": str(self.org_id) if self.org_id else None,
+            "session_id": str(self.session_id),
+            "timestamp": self.timestamp,
+            "request_id": str(self.request_id) if self.request_id else None,
+            "pipeline_run_id": str(self.pipeline_run_id) if self.pipeline_run_id else None,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class AuthFailureEvent:
+    """Event emitted on authentication failure.
+
+    Attributes:
+        reason: Failure reason code (e.g., "token_expired", "invalid_signature")
+        ip_address: Client IP address
+        user_agent: Client user agent string
+        timestamp: When failure occurred
+        request_id: Request correlation ID
+        user_id: User ID if known (e.g., from expired token)
+    """
+
+    reason: str
+    ip_address: str | None = None
+    user_agent: str | None = None
+    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
+    request_id: UUID | None = None
+    user_id: UUID | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for event emission."""
+        return {
+            "type": "auth.failure",
+            "reason": self.reason,
+            "ip_address": self.ip_address,
+            "user_agent": self.user_agent,
+            "timestamp": self.timestamp,
+            "request_id": str(self.request_id) if self.request_id else None,
+            "user_id": str(self.user_id) if self.user_id else None,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class TenantAccessDeniedEvent:
+    """Event emitted on cross-tenant access attempt.
+
+    Attributes:
+        user_org_id: User's organization ID
+        resource_org_id: Resource's organization ID
+        user_id: User who attempted access
+        timestamp: When violation occurred
+        request_id: Request correlation ID
+        pipeline_run_id: Pipeline run correlation ID
+    """
+
+    user_org_id: UUID
+    resource_org_id: UUID
+    user_id: UUID | None = None
+    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
+    request_id: UUID | None = None
+    pipeline_run_id: UUID | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for event emission."""
+        return {
+            "type": "tenant.access_denied",
+            "user_org_id": str(self.user_org_id),
+            "resource_org_id": str(self.resource_org_id),
+            "user_id": str(self.user_id) if self.user_id else None,
+            "timestamp": self.timestamp,
+            "request_id": str(self.request_id) if self.request_id else None,
+            "pipeline_run_id": str(self.pipeline_run_id) if self.pipeline_run_id else None,
+        }
+
+
+__all__ = [
+    "AuthFailureEvent",
+    "AuthLoginEvent",
+    "TenantAccessDeniedEvent",
+]