pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/main.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+from fastapi import FastAPI
+from starlette.middleware.cors import CORSMiddleware
+
+from pulse_engine.api.v1.router import v1_router
+from pulse_engine.config import Settings, get_settings
+from pulse_engine.core.error_handlers import register_exception_handlers
+from pulse_engine.core.security import (
+    CognitoTokenVerifier,
+    MockTokenVerifier,
+    TokenVerifier,
+)
+from pulse_engine.middleware.rate_limit import RateLimitMiddleware
+from pulse_engine.middleware.request_id import RequestIDMiddleware
+from pulse_engine.middleware.security_headers import SecurityHeadersMiddleware
+from pulse_engine.middleware.tenant import TenantMiddleware
+from pulse_engine.pipeline.router_modules import router as pipeline_modules_router
+from pulse_engine.pipeline.router_pipelines import router as pipeline_pipelines_router
+from pulse_engine.services.bootstrap import bootstrap_services
+
+if TYPE_CHECKING:
+    from pulse_engine.registry import ProductManifest
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+    settings: Settings = app.state.settings
+    manifest: ProductManifest | None = getattr(app.state, "manifest", None)
+
+    async with bootstrap_services(settings, manifest) as container:
+        app.state.opensearch = container.opensearch
+        app.state.knowledge_base = container.kb_service
+        app.state.db_session_factory = container.db_session_factory
+        app.state.orchestrator_adapter = container.orchestrator_adapter
+        app.state.celery_app = container.celery_app
+        app.state.pipeline = container.pipeline
+        app.state.s3_client = container.s3_client
+        app.state.token_issuer = container.token_issuer
+        app.state.pipeline_translators = container.pipeline_translators
+        app.state.pipeline_status_providers = container.pipeline_status_providers
+
+        # Start chain recovery polling if configured
+        recovery_task = None
+        if (
+            container.db_session_factory
+            and container.token_issuer
+            and settings.pulse_engine_url
+            and settings.pulse_job_token_secret
+        ):
+
+            async def _recovery_loop() -> None:
+                from pulse_engine.chain_recovery import ChainRecoveryTask
+                from pulse_engine.deployment.backend_deployment_repository import (
+                    BackendDeploymentRepository,
+                )
+                from pulse_engine.deployment.backends.registry import BackendRegistry
+                from pulse_engine.deployment.job_launcher import JobLauncher
+                from pulse_engine.deployment.repository import (
+                    RegistrationRepository,
+                )
+                from pulse_engine.extractor.repository import JobRepository
+                from pulse_engine.extractor.stage_repository import (
+                    StageRepository,
+                )
+
+                registry = BackendRegistry(settings)
+                factory = container.db_session_factory
+                assert factory is not None
+                while True:
+                    try:
+                        async with factory() as session:
+                            job_launcher = JobLauncher(
+                                registration_repo=RegistrationRepository(session),
+                                backend_deployment_repo=BackendDeploymentRepository(
+                                    session
+                                ),
+                                registry=registry,
+                                token_issuer=container.token_issuer,
+                                settings=settings,
+                            )
+                            task = ChainRecoveryTask(
+                                stage_repo=StageRepository(session),
+                                job_repo=JobRepository(session),
+                                job_launcher=job_launcher,
+                                orchestrator=container.orchestrator_adapter,
+                                settings=settings,
+                            )
+                            await task.check_once()
+                            await task.check_stalled_chains()
+                    except Exception:
+                        import structlog
+
+                        structlog.get_logger().warning(
+                            "chain_recovery_loop_error",
+                            exc_info=True,
+                        )
+                    await asyncio.sleep(60)
+
+            recovery_task = asyncio.create_task(_recovery_loop())
+
+        try:
+            yield
+        finally:
+            if recovery_task:
+                recovery_task.cancel()
+
+
+def _build_verifier(settings: Settings) -> TokenVerifier:
+    if settings.mock_jwt_secret:
+        return MockTokenVerifier(
+            secret=settings.mock_jwt_secret,
+            audience=settings.cognito_app_client_id,
+            issuer=settings.cognito_issuer,
+        )
+    return CognitoTokenVerifier(
+        jwks_url=settings.cognito_jwks_url,
+        issuer=settings.cognito_issuer,
+        audience=settings.cognito_app_client_id,
+    )
+
+
+def create_app(
+    settings: Settings | None = None,
+    token_verifier: TokenVerifier | None = None,
+    manifest: ProductManifest | None = None,
+) -> FastAPI:
+    if settings is None:
+        settings = get_settings()
+
+    if manifest:
+        from pulse_engine.registry import ManifestValidationError, validate_manifest
+
+        errors = validate_manifest(manifest)
+        if errors:
+            raise ManifestValidationError(errors)
+
+    app = FastAPI(
+        title="Pulse Platform",
+        version=settings.app_version,
+        lifespan=lifespan,
+        docs_url=None if settings.app_env == "production" else "/docs",
+        redoc_url=None if settings.app_env == "production" else "/redoc",
+        openapi_url=None if settings.app_env == "production" else "/openapi.json",
+    )
+
+    app.state.settings = settings
+    app.state.manifest = manifest
+
+    register_exception_handlers(app)
+
+    verifier = token_verifier or _build_verifier(settings)
+
+    # Middleware is applied in reverse order — outermost first
+    app.add_middleware(
+        TenantMiddleware,
+        verifier=verifier,
+        job_token_secret=settings.pulse_job_token_secret,
+    )
+    app.add_middleware(SecurityHeadersMiddleware)
+    app.add_middleware(RequestIDMiddleware)
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=(
+            ["*"]
+            if settings.app_env == "development"
+            else [settings.pulse_engine_url]
+            if settings.pulse_engine_url
+            else []
+        ),
+        allow_credentials=True,
+        allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
+        allow_headers=["Authorization", "Content-Type", "X-Request-ID"],
+        expose_headers=["X-Request-ID"],
+    )
+    app.add_middleware(
+        RateLimitMiddleware,
+        max_requests=100,
+        window_seconds=60,
+    )
+
+    app.include_router(v1_router)
+    app.include_router(pipeline_modules_router)
+    app.include_router(pipeline_pipelines_router)
+
+    if manifest:
+        for router in manifest.routers:
+            app.include_router(router, prefix="/api/v1")
+
+    return app

pulse_engine/mcp/__main__.py

@@ -0,0 +1,5 @@
+import asyncio
+
+from pulse_engine.mcp.server import run_mcp_server
+
+asyncio.run(run_mcp_server())

pulse_engine/mcp/server.py

@@ -0,0 +1,108 @@
+"""MCP server for Pulse Platform.
+
+Exposes KB, Jobs, Processor, Pipeline, and Module tools.
+"""
+
+from __future__ import annotations
+
+import functools
+import importlib
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from mcp.server.fastmcp import FastMCP
+from pydantic import ValidationError
+
+from pulse_engine.config import get_settings as _get_settings
+from pulse_engine.core.exceptions import (
+    AppError,
+    NotFoundError,
+    ServiceUnavailableError,
+    TooManyRequestsError,
+)
+from pulse_engine.services.bootstrap import ServiceContainer
+
+_s = _get_settings()
+mcp = FastMCP("Pulse Platform", host=_s.mcp_sse_host, port=_s.mcp_sse_port)
+
+_container: ServiceContainer | None = None
+
+
+def get_container() -> ServiceContainer:
+    if _container is None:
+        raise RuntimeError("MCP server not initialized — no service container")
+    return _container
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def handle_mcp_errors(fn: F) -> F:
+    @functools.wraps(fn)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return await fn(*args, **kwargs)
+        except NotFoundError as e:
+            return {"error": True, "error_type": "NotFoundError", "message": e.message}
+        except TooManyRequestsError as e:
+            return {
+                "error": True,
+                "error_type": "TooManyRequestsError",
+                "message": e.message,
+            }
+        except ServiceUnavailableError as e:
+            return {
+                "error": True,
+                "error_type": "ServiceUnavailableError",
+                "message": e.message,
+            }
+        except ValidationError as e:
+            return {
+                "error": True,
+                "error_type": "ValidationError",
+                "message": str(e),
+            }
+        except AppError as e:
+            return {
+                "error": True,
+                "error_type": type(e).__name__,
+                "message": e.message,
+            }
+
+    return wrapper  # type: ignore[return-value]
+
+
+def create_mcp_server(
+    manifest: Any | None = None,
+) -> FastMCP:
+    """Return the module-level FastMCP instance, optionally loading product tools."""
+    if manifest:
+        for module_path in manifest.mcp_tool_modules:
+            importlib.import_module(module_path)
+    return mcp
+
+
+async def run_mcp_server(manifest: Any | None = None) -> None:
+    global _container
+
+    from pulse_engine.config import get_settings
+    from pulse_engine.services.bootstrap import bootstrap_services
+
+    settings = get_settings()
+
+    async with bootstrap_services(settings, manifest) as container:
+        _container = container
+
+        # Import engine tool modules to trigger @mcp.tool() registrations
+        import pulse_engine.mcp.tools_jobs  # noqa: F401
+        import pulse_engine.mcp.tools_kb  # noqa: F401
+        import pulse_engine.mcp.tools_modules  # noqa: F401
+        import pulse_engine.mcp.tools_pipelines  # noqa: F401
+        import pulse_engine.mcp.tools_processor  # noqa: F401
+
+        # Import product tool modules
+        if manifest:
+            for module_path in manifest.mcp_tool_modules:
+                importlib.import_module(module_path)
+
+        await mcp.run_sse_async()

pulse_engine/mcp/tools_jobs.py

@@ -0,0 +1,159 @@
+"""MCP tools for Job operations."""
+
+from typing import Any
+
+from pulse_engine.core.exceptions import ServiceUnavailableError
+from pulse_engine.extractor.repository import JobRepository
+from pulse_engine.extractor.schemas import (
+    CreateJobRequest,
+    JobPriority,
+    StatusUpdateRequest,
+)
+from pulse_engine.extractor.service import JobService
+from pulse_engine.mcp.server import get_container, handle_mcp_errors, mcp
+
+
+def _build_job_service(
+    session: Any,
+    container: Any,
+) -> JobService:
+    repo = JobRepository(session)
+    return JobService(
+        repository=repo,
+        orchestrator=container.orchestrator_adapter,
+        settings=container.settings,
+    )
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_register(
+    tenant_id: str,
+    job_type: str,
+    product: str,
+    priority: str = "normal",
+    orchestrator_run_id: str | None = None,
+    parameters: dict[str, Any] | None = None,
+    callback_url: str | None = None,
+) -> dict[str, Any]:
+    """Register a new extraction job for a tenant."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        request = CreateJobRequest(
+            job_type=job_type,
+            product=product,
+            priority=JobPriority(priority),
+            orchestrator_run_id=orchestrator_run_id,
+            parameters=parameters or {},
+            callback_url=callback_url,
+        )
+        result = await service.register_job(tenant_id, request)
+        return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_get(
+    tenant_id: str,
+    job_id: str,
+) -> dict[str, Any]:
+    """Get details of a specific job."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        result = await service.get_job(tenant_id, job_id)
+        return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_list(
+    tenant_id: str,
+    status: str | None = None,
+    product: str | None = None,
+    job_type: str | None = None,
+    limit: int = 20,
+    offset: int = 0,
+    sort_by: str = "created_at",
+    order: str = "desc",
+) -> dict[str, Any]:
+    """List jobs for a tenant with optional filters."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        result = await service.list_jobs(
+            tenant_id,
+            status=status,
+            product=product,
+            job_type=job_type,
+            limit=limit,
+            offset=offset,
+            sort_by=sort_by,
+            order=order,
+        )
+        return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_push_status(
+    tenant_id: str,
+    job_id: str,
+    status: str,
+    result_summary: dict[str, Any] | None = None,
+    error: str | None = None,
+) -> dict[str, Any]:
+    """Push a status update for a job."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        from pulse_engine.extractor.schemas import JobStatus
+
+        request = StatusUpdateRequest(
+            status=JobStatus(status),
+            result_summary=result_summary,
+            error=error,
+        )
+        result = await service.push_status(tenant_id, job_id, request)
+        return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_cancel(
+    tenant_id: str,
+    job_id: str,
+) -> dict[str, Any]:
+    """Cancel a running or pending job."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        result = await service.cancel_job(tenant_id, job_id)
+        return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def jobs_delete(
+    tenant_id: str,
+    job_id: str,
+) -> dict[str, Any]:
+    """Delete a job record."""
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+    async with container.db_session_factory() as session:
+        service = _build_job_service(session, container)
+        deleted = await service.delete_job(tenant_id, job_id)
+        return {"deleted": deleted, "job_id": job_id}

pulse_engine/mcp/tools_kb.py

@@ -0,0 +1,88 @@
+"""MCP tools for Knowledge Base operations."""
+
+from typing import Any
+
+from pulse_engine.mcp.server import get_container, handle_mcp_errors, mcp
+from pulse_engine.storage.schemas import Document, SearchQuery, SearchQueryType
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_store_documents(
+    tenant_id: str,
+    documents: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """Store documents in the Knowledge Base for a tenant."""
+    container = get_container()
+    docs = [Document(**d) for d in documents]
+    result = await container.kb_service.store_documents(tenant_id, docs)
+    return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_retrieve_document(
+    tenant_id: str,
+    document_id: str,
+) -> dict[str, Any]:
+    """Retrieve a single document by ID from the Knowledge Base."""
+    container = get_container()
+    doc = await container.kb_service.retrieve_document(tenant_id, document_id)
+    return doc.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_search(
+    tenant_id: str,
+    query_text: str,
+    query_type: str = "text",
+    filters: dict[str, Any] | None = None,
+    size: int = 10,
+) -> dict[str, Any]:
+    """Search documents in the Knowledge Base."""
+    container = get_container()
+    query = SearchQuery(
+        query_type=SearchQueryType(query_type),
+        query_text=query_text,
+        filters=filters or {},
+        size=size,
+    )
+    result = await container.kb_service.search(tenant_id, query)
+    return result.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_delete_document(
+    tenant_id: str,
+    document_id: str,
+) -> dict[str, Any]:
+    """Delete a document from the Knowledge Base."""
+    container = get_container()
+    await container.kb_service.delete_document(tenant_id, document_id)
+    return {"deleted": True, "document_id": document_id}
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_get_stats(
+    tenant_id: str,
+) -> dict[str, Any]:
+    """Get Knowledge Base statistics for a tenant."""
+    container = get_container()
+    stats = await container.kb_service.get_stats(tenant_id)
+    return stats.model_dump(mode="json")
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def kb_run_query(
+    tenant_id: str,
+    sql: str,
+    parameters: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Execute an Athena SQL query against the Knowledge Base."""
+    container = get_container()
+    result = await container.kb_service.execute_query(tenant_id, sql, parameters)
+    return result.model_dump(mode="json")

pulse_engine/mcp/tools_modules.py

@@ -0,0 +1,115 @@
+"""MCP tools for Module Registry operations."""
+
+from typing import Any
+
+from pulse_engine.core.exceptions import NotFoundError, ServiceUnavailableError
+from pulse_engine.mcp.server import get_container, handle_mcp_errors, mcp
+from pulse_engine.pipeline.repositories import ModuleRegistryRepository
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def modules_register(
+    tenant_id: str,
+    product: str,
+    modules: dict[str, dict[str, str]],
+) -> dict[str, Any]:
+    """Register pipeline module images for a product.
+
+    Args:
+        tenant_id: Tenant identifier.
+        product: Product name (e.g. "sample-product").
+        modules: Map of module name to {"image": "<ecr-uri>:<tag>"}.
+            Example: {"extractor": {"image": "123.ecr/extractor:latest"}}
+
+    Returns:
+        {"product": "...", "modules": {"extractor": {"status": "registered"}, ...}}
+    """
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+
+    async with container.db_session_factory() as session:
+        repo = ModuleRegistryRepository(session)
+        result_modules: dict[str, dict[str, str]] = {}
+
+        for module_name, module_def in modules.items():
+            image = module_def.get("image", "")
+            if not image:
+                result_modules[module_name] = {
+                    "status": "error",
+                    "detail": "missing image",
+                }
+                continue
+            await repo.upsert(tenant_id, product, module_name, image)
+            result_modules[module_name] = {"status": "registered"}
+
+        return {"product": product, "modules": result_modules}
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def modules_list(
+    tenant_id: str,
+    product: str,
+) -> dict[str, Any]:
+    """List registered modules for a product.
+
+    Args:
+        tenant_id: Tenant identifier.
+        product: Product name.
+
+    Returns:
+        {"product": "...", "modules": [{"name": "...", "image": "..."}, ...]}
+    """
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+
+    async with container.db_session_factory() as session:
+        repo = ModuleRegistryRepository(session)
+        entries = await repo.list_by_product(tenant_id, product)
+        return {
+            "product": product,
+            "modules": [
+                {
+                    "name": e.module_name,
+                    "image": e.image,
+                    "updated_at": e.updated_at.isoformat() if e.updated_at else None,
+                }
+                for e in entries
+            ],
+        }
+
+
+@mcp.tool()
+@handle_mcp_errors
+async def modules_delete(
+    tenant_id: str,
+    product: str,
+    module_name: str,
+) -> dict[str, Any]:
+    """Delete a module registration.
+
+    Args:
+        tenant_id: Tenant identifier.
+        product: Product name.
+        module_name: Module to delete.
+
+    Returns:
+        {"deleted": true, "product": "...", "module_name": "..."}
+    """
+    container = get_container()
+    if container.db_session_factory is None:
+        raise ServiceUnavailableError("Database not configured")
+
+    async with container.db_session_factory() as session:
+        repo = ModuleRegistryRepository(session)
+        deleted = await repo.delete(tenant_id, product, module_name)
+        if not deleted:
+            raise NotFoundError(
+                f"Module '{module_name}' not found for product '{product}'",
+                product=product,
+                module_name=module_name,
+            )
+        return {"deleted": True, "product": product, "module_name": module_name}