cloud-dog-api-kit 0.13.0__py3-none-any.whl

cloud_dog_api_kit/lifecycle/shutdown.py

@@ -0,0 +1,178 @@
+# Copyright 2026 Cloud-Dog, Viewdeck Engineering Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# cloud_dog_api_kit — Graceful shutdown management
+#
+# Licence: Proprietary — Cloud-Dog AI Platform
+# Owner: Cloud-Dog AI
+# Description: In-flight request draining and signal-handler integration for
+#   deterministic graceful shutdown.
+# Related requirements: FR18.9
+# Related architecture: SA1
+
+"""Graceful shutdown support for cloud_dog_api_kit."""
+
+from __future__ import annotations
+
+import asyncio
+import signal
+import threading
+from typing import Any, Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+
+from cloud_dog_api_kit.envelopes.error import error_envelope
+
+
+class GracefulShutdownManager:
+    """Track in-flight requests and coordinate graceful shutdown.
+
+    Args:
+        drain_timeout_seconds: Maximum wait time for in-flight requests.
+
+    Related tests: UT1.45_GracefulShutdown, ST1.13_StartupLifecycleIntegration
+    """
+
+    def __init__(self, drain_timeout_seconds: float = 5.0) -> None:
+        if drain_timeout_seconds < 0:
+            raise ValueError("drain_timeout_seconds must be >= 0")
+        self._drain_timeout_seconds = drain_timeout_seconds
+        self._lock = threading.Lock()
+        self._active_requests = 0
+        self._shutting_down = False
+        self._all_requests_drained = asyncio.Event()
+        self._all_requests_drained.set()
+
+    @property
+    def active_requests(self) -> int:
+        """Current number of active requests."""
+        with self._lock:
+            return self._active_requests
+
+    @property
+    def shutting_down(self) -> bool:
+        """Whether shutdown has started."""
+        with self._lock:
+            return self._shutting_down
+
+    def mark_request_started(self) -> bool:
+        """Mark a request start.
+
+        Returns:
+            True when request is accepted, False when shutdown is active.
+        """
+        with self._lock:
+            if self._shutting_down:
+                return False
+            self._active_requests += 1
+            self._all_requests_drained.clear()
+            return True
+
+    def mark_request_finished(self) -> None:
+        """Mark request completion."""
+        with self._lock:
+            if self._active_requests > 0:
+                self._active_requests -= 1
+            if self._active_requests == 0:
+                self._all_requests_drained.set()
+
+    def set_shutting_down(self) -> None:
+        """Set shutdown flag to reject new requests."""
+        with self._lock:
+            self._shutting_down = True
+            if self._active_requests == 0:
+                self._all_requests_drained.set()
+
+    async def drain(self) -> bool:
+        """Wait for in-flight requests to drain.
+
+        Returns:
+            True if drained before timeout, otherwise False.
+        """
+        if self.active_requests == 0:
+            return True
+        try:
+            await asyncio.wait_for(self._all_requests_drained.wait(), timeout=self._drain_timeout_seconds)
+            return True
+        except asyncio.TimeoutError:
+            return False
+
+    async def initiate_shutdown(self) -> bool:
+        """Start shutdown and drain requests."""
+        self.set_shutting_down()
+        return await self.drain()
+
+
+class ShutdownDrainMiddleware(BaseHTTPMiddleware):
+    """Reject new requests during shutdown and track in-flight requests.
+
+    Args:
+        app: The ASGI application.
+        manager: Shared graceful shutdown manager.
+    """
+
+    def __init__(self, app: Any, manager: GracefulShutdownManager) -> None:
+        super().__init__(app)
+        self._manager = manager
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        """Track active request and reject when shutting down."""
+        if not self._manager.mark_request_started():
+            request_id = getattr(request.state, "request_id", "")
+            correlation_id = getattr(request.state, "correlation_id", None)
+            return JSONResponse(
+                status_code=503,
+                content=error_envelope(
+                    code="INTERNAL_ERROR",
+                    message="Service is shutting down",
+                    details=None,
+                    retryable=True,
+                    request_id=request_id,
+                    correlation_id=correlation_id,
+                ),
+            )
+        try:
+            return await call_next(request)
+        finally:
+            self._manager.mark_request_finished()
+
+
+def install_shutdown_signal_handlers(
+    manager: GracefulShutdownManager,
+    loop: asyncio.AbstractEventLoop | None = None,
+) -> dict[int, bool]:
+    """Attempt to install SIGTERM/SIGINT handlers for graceful shutdown.
+
+    Returns:
+        Mapping of signal number to whether installation succeeded.
+    """
+    try:
+        active_loop = loop or asyncio.get_running_loop()
+    except RuntimeError:
+        return {signal.SIGTERM: False, signal.SIGINT: False}
+
+    results: dict[int, bool] = {}
+
+    def _schedule_shutdown() -> None:
+        active_loop.create_task(manager.initiate_shutdown())
+
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        try:
+            active_loop.add_signal_handler(sig, _schedule_shutdown)
+            results[sig] = True
+        except (NotImplementedError, RuntimeError, ValueError):
+            results[sig] = False
+    return results

cloud_dog_api_kit/mcp/__init__.py

@@ -0,0 +1,122 @@
+# Copyright 2026 Cloud-Dog, Viewdeck Engineering Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# cloud_dog_api_kit — MCP gateway
+#
+# Licence: Proprietary — Cloud-Dog AI Platform
+# Owner: Cloud-Dog AI
+# Description: REST-to-MCP gateway helpers.
+# Related requirements: FR14.1
+# Related architecture: SA1
+
+"""MCP gateway helpers for cloud_dog_api_kit."""
+
+from __future__ import annotations
+
+from cloud_dog_api_kit.mcp.client_transport import (
+    HTTPJSONRPCConfig,
+    HTTPJSONRPCTransport,
+    LegacySSETransport,
+    MCPProtocolError,
+    MCPSessionError,
+    MCPTransport,
+    MCPTransportError,
+    StdioConfig,
+    StdioTransport,
+    StreamableHTTPConfig,
+    StreamableHTTPTransport,
+)
+from cloud_dog_api_kit.mcp.async_jobs import AsyncJobStore, InMemoryAsyncJobStore
+from cloud_dog_api_kit.mcp.gateway import MCPToolDefinition, create_mcp_tool_from_endpoint
+from cloud_dog_api_kit.mcp.contract import MCPContractRegistration, register_mcp_contract
+from cloud_dog_api_kit.mcp.error_mapper import map_legacy_mcp_payload
+from cloud_dog_api_kit.mcp.legacy_sse import LegacySSEConfig
+from cloud_dog_api_kit.mcp.session import SESSION_HEADER, McpSession, McpSessionManager
+from cloud_dog_api_kit.mcp.sync_handler import (
+    BUDGET_EXCEEDED_CODE,
+    dispatch_with_sync_class,
+    format_a2a_task_response,
+    format_async_rest_response,
+    format_budget_exceeded_mcp_error,
+    validate_blocking_siblings,
+)
+from cloud_dog_api_kit.mcp.client_sdk import (
+    AsyncJobClient,
+    JobFailedError,
+    JobTimeoutError,
+)
+from cloud_dog_api_kit.mcp.tool_router import (
+    DEFAULT_SYNC_BUDGET_SECONDS,
+    MAX_SYNC_BUDGET_SECONDS,
+    SYNC_CLASS_ASYNC,
+    SYNC_CLASS_DEFAULT,
+    SYNC_CLASS_PROGRESS,
+    SYNC_CLASSES,
+    ToolContract,
+    register_tool_router,
+)
+from cloud_dog_api_kit.mcp.transport import (
+    AsyncJobResultShape,
+    MCPResource,
+    MCPResourceContent,
+    ResourcesHandler,
+    register_mcp_routes,
+)
+
+__all__ = [
+    "AsyncJobStore",
+    "InMemoryAsyncJobStore",
+    "MCPContractRegistration",
+    "MCPToolDefinition",
+    "MCPProtocolError",
+    "MCPSessionError",
+    "MCPTransport",
+    "MCPTransportError",
+    "McpSession",
+    "McpSessionManager",
+    "SESSION_HEADER",
+    "HTTPJSONRPCConfig",
+    "HTTPJSONRPCTransport",
+    "LegacySSEConfig",
+    "LegacySSETransport",
+    "AsyncJobResultShape",
+    "MCPResource",
+    "MCPResourceContent",
+    "ResourcesHandler",
+    "StdioConfig",
+    "StdioTransport",
+    "StreamableHTTPConfig",
+    "StreamableHTTPTransport",
+    "BUDGET_EXCEEDED_CODE",
+    "DEFAULT_SYNC_BUDGET_SECONDS",
+    "MAX_SYNC_BUDGET_SECONDS",
+    "SYNC_CLASS_ASYNC",
+    "SYNC_CLASS_DEFAULT",
+    "SYNC_CLASS_PROGRESS",
+    "SYNC_CLASSES",
+    "ToolContract",
+    "dispatch_with_sync_class",
+    "format_a2a_task_response",
+    "format_async_rest_response",
+    "format_budget_exceeded_mcp_error",
+    "validate_blocking_siblings",
+    "create_mcp_tool_from_endpoint",
+    "map_legacy_mcp_payload",
+    "register_mcp_contract",
+    "register_mcp_routes",
+    "register_tool_router",
+    "AsyncJobClient",
+    "JobFailedError",
+    "JobTimeoutError",
+]

cloud_dog_api_kit/mcp/async_jobs.py

@@ -0,0 +1,126 @@
+# Copyright 2026 Cloud-Dog, Viewdeck Engineering Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# cloud_dog_api_kit — MCP async job helpers
+#
+# Licence: Proprietary — Cloud-Dog AI Platform
+# Owner: Cloud-Dog AI
+# Description: Shared async job-store protocol and in-memory reference
+#   implementation for wait=false MCP tool calls.
+# Related requirements: FR18.1
+# Related architecture: SA1
+
+"""Async MCP tool-call job helpers."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import threading
+import uuid
+from collections.abc import Awaitable, Callable
+from typing import Any, Protocol
+
+AsyncJobContext = dict[str, Any]
+AsyncJobStatus = dict[str, Any]
+
+
+class AsyncJobStore(Protocol):
+    """Protocol for async MCP tool-call job stores."""
+
+    def submit(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        context: AsyncJobContext,
+    ) -> str | Awaitable[str]:
+        """Submit a tool call for async execution and return a job id."""
+
+    def get_status(self, job_id: str) -> AsyncJobStatus | Awaitable[AsyncJobStatus]:
+        """Return the current status payload for a submitted job."""
+
+
+class InMemoryAsyncJobStore:
+    """In-memory async job store for local tests and simple service deployments."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._jobs: dict[str, AsyncJobStatus] = {}
+        self._tasks: dict[str, asyncio.Task[Any]] = {}
+
+    def submit(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        context: AsyncJobContext,
+    ) -> str:
+        """Submit a tool call for background execution."""
+        runner = context.get("runner")
+        if not callable(runner):
+            raise TypeError("async job context requires a callable 'runner'")
+
+        result_formatter = context.get("result_formatter")
+        if result_formatter is not None and not callable(result_formatter):
+            raise TypeError("async job context 'result_formatter' must be callable when provided")
+
+        job_id = f"job-{uuid.uuid4().hex}"
+        with self._lock:
+            self._jobs[job_id] = {"status": "pending", "tool_name": tool_name, "arguments": dict(arguments)}
+
+        task = asyncio.create_task(self._run_job(job_id, runner, result_formatter))
+        with self._lock:
+            self._tasks[job_id] = task
+        return job_id
+
+    async def _run_job(
+        self,
+        job_id: str,
+        runner: Callable[[], Any],
+        result_formatter: Callable[[Any], Any] | None,
+    ) -> None:
+        """Execute a submitted async job and persist its lifecycle state."""
+        with self._lock:
+            job = dict(self._jobs.get(job_id) or {})
+            job["status"] = "running"
+            self._jobs[job_id] = job
+
+        try:
+            result = runner()
+            if inspect.isawaitable(result):
+                result = await result
+            if result_formatter is not None:
+                result = result_formatter(result)
+            with self._lock:
+                job = dict(self._jobs.get(job_id) or {})
+                job["status"] = "completed"
+                job["result"] = result
+                job.pop("error", None)
+                self._jobs[job_id] = job
+        except Exception as exc:
+            with self._lock:
+                job = dict(self._jobs.get(job_id) or {})
+                job["status"] = "failed"
+                job["error"] = str(exc)
+                self._jobs[job_id] = job
+        finally:
+            with self._lock:
+                self._tasks.pop(job_id, None)
+
+    def get_status(self, job_id: str) -> AsyncJobStatus:
+        """Return the current status payload for a job id."""
+        with self._lock:
+            status = self._jobs.get(job_id)
+            if status is None:
+                return {"status": "not_found", "error": "Job not found"}
+            return dict(status)

cloud_dog_api_kit/mcp/client_sdk.py

@@ -0,0 +1,235 @@
+# Copyright 2026 Cloud-Dog, Viewdeck Engineering Limited
+# SPDX-License-Identifier: Apache-2.0
+"""PS-95 client SDK helpers for three-mode tool invocation (W28D-309).
+
+Provides:
+- submit(): fire-and-forget async job submission (Mode 3)
+- wait(): poll until job completes or timeout
+- call_with_auto_wait(): smart caller that handles all three modes
+- progress_token_handler(): handles Mode 2 progress notifications
+- retry_idempotent(): idempotent retry with automatic dedup
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Callable, Awaitable
+
+# PS-95 constants (consistent with sync_handler.py)
+DEFAULT_POLL_INTERVAL_SECONDS = 1.0
+DEFAULT_CLIENT_TIMEOUT_SECONDS = 300.0  # 5 minutes (PS-95 §9 worker_timeout)
+BUDGET_EXCEEDED_CODE = -32000
+
+
+class AsyncJobClient:
+    """Client-side helper for PS-95 three-mode tool calls.
+
+    Wraps a transport (MCP/REST/A2A) and provides high-level helpers
+    that handle the three sync_class modes transparently.
+    """
+
+    def __init__(
+        self,
+        *,
+        call_tool: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]],
+        poll_job: Callable[[str], Awaitable[dict[str, Any]]],
+        poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
+        timeout: float = DEFAULT_CLIENT_TIMEOUT_SECONDS,
+    ) -> None:
+        """Initialize.
+
+        Args:
+            call_tool: async callable(tool_name, arguments) -> server response dict
+            poll_job: async callable(job_id) -> job status dict
+            poll_interval: seconds between poll requests
+            timeout: max seconds to wait for job completion
+        """
+        self._call_tool = call_tool
+        self._poll_job = poll_job
+        self._poll_interval = poll_interval
+        self._timeout = timeout
+
+    async def submit(self, tool_name: str, arguments: dict[str, Any]) -> str:
+        """Submit a tool call and return the job_id immediately (Mode 3 client).
+
+        If the server returns a sync result, wraps it as a completed job.
+        """
+        response = await self._call_tool(tool_name, arguments)
+        job_id = _extract_job_id(response)
+        if job_id is None:
+            raise ValueError(
+                f"Server did not return a job_id for tool {tool_name!r}. "
+                "Use call_with_auto_wait() for tools that may return sync results."
+            )
+        return job_id
+
+    async def wait(self, job_id: str, *, timeout: float | None = None) -> dict[str, Any]:
+        """Poll a job until completion or timeout.
+
+        Returns the final job status dict with 'result' on success.
+        Raises TimeoutError with the job_id if timeout expires.
+        """
+        deadline = time.monotonic() + (timeout or self._timeout)
+        while True:
+            status = await self._poll_job(job_id)
+            state = str(status.get("status", "")).lower()
+            if state in ("completed", "succeeded"):
+                return status
+            if state in ("failed", "cancelled", "timeout", "dead_lettered"):
+                raise JobFailedError(job_id=job_id, status=status)
+            if time.monotonic() >= deadline:
+                raise JobTimeoutError(job_id=job_id, last_status=status)
+            await asyncio.sleep(self._poll_interval)
+
+    async def call_with_auto_wait(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        *,
+        timeout: float | None = None,
+        on_progress: Callable[[float, float | None, str | None], Awaitable[None]] | None = None,
+    ) -> dict[str, Any]:
+        """Call a tool with automatic mode detection and waiting.
+
+        - If server returns sync result (Mode 1 completed): returns immediately.
+        - If server returns job_id (Mode 3 or budget-exceeded): polls until done.
+        - If server sends progress (Mode 2): calls on_progress, then returns result.
+
+        Returns dict with 'result' key on success.
+        Raises JobTimeoutError with job_id on timeout (client can resume polling).
+        """
+        response = await self._call_tool(tool_name, arguments)
+
+        # Check for sync result (Mode 1 completed within budget)
+        if _is_sync_result(response):
+            return {"status": "completed", "result": response.get("result", response)}
+
+        # Check for budget-exceeded error (Mode 1 timeout)
+        if _is_budget_exceeded(response):
+            job_id = _extract_job_id_from_error(response)
+            if job_id:
+                return await self.wait(job_id, timeout=timeout)
+            raise ValueError("Budget exceeded but no job_id in error response")
+
+        # Async/job_id response (Mode 3)
+        job_id = _extract_job_id(response)
+        if job_id:
+            return await self.wait(job_id, timeout=timeout)
+
+        # Fallback: treat entire response as sync result
+        return {"status": "completed", "result": response}
+
+    async def retry_idempotent(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        *,
+        idempotency_key: str | None = None,
+        max_retries: int = 3,
+        timeout: float | None = None,
+    ) -> dict[str, Any]:
+        """Retry a tool call with idempotency guarantees.
+
+        If the call returns a job_id, subsequent retries poll the existing job
+        instead of resubmitting (deduplication).
+        """
+        last_error: Exception | None = None
+        known_job_id: str | None = None
+
+        for attempt in range(max_retries):
+            try:
+                if known_job_id:
+                    return await self.wait(known_job_id, timeout=timeout)
+
+                args = dict(arguments)
+                if idempotency_key:
+                    args["_idempotency_key"] = idempotency_key
+
+                response = await self._call_tool(tool_name, args)
+                job_id = _extract_job_id(response)
+                if job_id:
+                    known_job_id = job_id
+                    return await self.wait(job_id, timeout=timeout)
+
+                if _is_sync_result(response):
+                    return {"status": "completed", "result": response.get("result", response)}
+
+                return {"status": "completed", "result": response}
+
+            except JobTimeoutError as e:
+                known_job_id = e.job_id
+                last_error = e
+            except Exception as e:
+                last_error = e
+
+            if attempt < max_retries - 1:
+                await asyncio.sleep(self._poll_interval * (attempt + 1))
+
+        raise last_error or RuntimeError("Retry exhausted with no result or error")
+
+
+# ── Response parsing helpers ──
+
+
+def _extract_job_id(response: dict[str, Any]) -> str | None:
+    """Extract job_id from various server response formats."""
+    if "job_id" in response:
+        return str(response["job_id"])
+    if "content" in response and isinstance(response["content"], list):
+        for item in response["content"]:
+            if isinstance(item, dict) and "job_id" in item:
+                return str(item["job_id"])
+    return None
+
+
+def _extract_job_id_from_error(response: dict[str, Any]) -> str | None:
+    """Extract job_id from budget-exceeded error response."""
+    error = response.get("error", {})
+    if isinstance(error, dict):
+        data = error.get("data", {})
+        if isinstance(data, dict) and "job_id" in data:
+            return str(data["job_id"])
+    return None
+
+
+def _is_sync_result(response: dict[str, Any]) -> bool:
+    """True if the response is a completed sync result (not a job envelope)."""
+    if "job_id" in response:
+        return False
+    if "error" in response:
+        return False
+    if response.get("mode") == "async":
+        return False
+    if response.get("status") == "submitted":
+        return False
+    return "result" in response or "content" in response
+
+
+def _is_budget_exceeded(response: dict[str, Any]) -> bool:
+    """True if the response is a PS-95 budget-exceeded error."""
+    error = response.get("error", {})
+    if isinstance(error, dict):
+        return error.get("code") == BUDGET_EXCEEDED_CODE
+    return False
+
+
+# ── Exceptions ──
+
+
+class JobTimeoutError(Exception):
+    """Raised when polling exceeds timeout. Carries job_id for resumption."""
+
+    def __init__(self, job_id: str, last_status: dict[str, Any] | None = None) -> None:
+        self.job_id = job_id
+        self.last_status = last_status
+        super().__init__(f"Job {job_id} timed out waiting for completion")
+
+
+class JobFailedError(Exception):
+    """Raised when a polled job reaches a terminal failure state."""
+
+    def __init__(self, job_id: str, status: dict[str, Any]) -> None:
+        self.job_id = job_id
+        self.status = status
+        super().__init__(f"Job {job_id} failed: {status.get('error', 'unknown')}")