codrspot-processor-mcp 0.1.0__py3-none-any.whl

codepreproc_client/layer1_business/mcp/lifecycle.py

@@ -0,0 +1,649 @@
+"""MCP connection lifecycle logging and instrumented server/session wrappers."""
+
+from __future__ import annotations
+
+import json
+import logging
+import signal
+import threading
+from contextlib import AsyncExitStack
+import atexit
+import time
+from itertools import count
+from typing import Any
+
+import anyio
+import mcp.types as types
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.server.experimental.request_context import Experimental
+from mcp.server.lowlevel.server import Server
+from mcp.server.lowlevel.server import logger as sdk_server_logger
+from mcp.server.lowlevel.server import request_ctx
+from mcp.server.models import InitializationOptions
+from mcp.server.session import InitializationState, ServerSession
+from mcp.shared.context import RequestContext
+from mcp.shared.exceptions import McpError
+from mcp.shared.message import ServerMessageMetadata, SessionMessage
+from mcp.shared.session import RequestResponder
+from contracts.dtos import MCPConnectionVerbosity, WorldModelLease
+
+_world_model_lease_logger = logging.getLogger("codepreproc_client.mcp.world_model_lease")
+
+
+class WorldModelLeaseGuard:
+    """Owns the temporary world-model copy for a session and guarantees deletion.
+
+    Architecture doc S5/S12: the leased world model is kept only in
+    session memory/scratch and MUST be discarded on the client's last MCP
+    call - not persisted to disk, not reused offline, deleted even on
+    abrupt termination. This registers teardown three ways so no exit
+    path skips it: explicit `discard()`/context-manager exit, `atexit`,
+    and the process signal handlers (SIGINT/SIGTERM; SIGBREAK on Windows).
+    """
+
+    _instances_lock = threading.Lock()
+    _instances: "set[WorldModelLeaseGuard]" = set()
+    _signal_handlers_installed = False
+
+    def __init__(self) -> None:
+        self._lease: WorldModelLease | None = None
+        with WorldModelLeaseGuard._instances_lock:
+            WorldModelLeaseGuard._instances.add(self)
+            WorldModelLeaseGuard._ensure_global_teardown_hooks()
+
+    @property
+    def lease(self) -> WorldModelLease | None:
+        return self._lease
+
+    def set_lease(self, lease: WorldModelLease) -> None:
+        self._lease = lease
+
+    def discard(self) -> None:
+        """Drop the in-memory temporary world-model copy. Idempotent."""
+
+        if self._lease is not None:
+            _world_model_lease_logger.debug("world_model_lease_discarded lease_id=%s", self._lease.lease_id)
+        self._lease = None
+
+    def __enter__(self) -> "WorldModelLeaseGuard":
+        return self
+
+    def __exit__(self, *exc_info: Any) -> None:
+        self.discard()
+
+    @classmethod
+    def _ensure_global_teardown_hooks(cls) -> None:
+        if cls._signal_handlers_installed:
+            return
+        cls._signal_handlers_installed = True
+        atexit.register(cls._discard_all)
+        for sig_name in ("SIGINT", "SIGTERM", "SIGBREAK"):
+            sig = getattr(signal, sig_name, None)
+            if sig is None:
+                continue
+            try:
+                previous_handler = signal.getsignal(sig)
+            except (ValueError, OSError):
+                # Not in the main thread, or unsupported on this platform - atexit still covers this.
+                continue
+
+            def _handler(signum: int, frame: Any, _previous: Any = previous_handler) -> None:
+                cls._discard_all()
+                if callable(_previous):
+                    _previous(signum, frame)
+
+            try:
+                signal.signal(sig, _handler)
+            except (ValueError, OSError):
+                continue
+
+    @classmethod
+    def _discard_all(cls) -> None:
+        with cls._instances_lock:
+            for instance in cls._instances:
+                instance.discard()
+
+
+class MCPConnectionLifecycleLogger:
+    """Log MCP transport/session lifecycle events to file and client notifications."""
+
+    def __init__(self, logger: logging.Logger, verbosity: MCPConnectionVerbosity = "basic") -> None:
+        self.logger = logger
+        self.verbosity = verbosity
+        self._session_counter = count(1)
+
+    @property
+    def enabled(self) -> bool:
+        return self.verbosity != "off"
+
+    @property
+    def verbose(self) -> bool:
+        return self.verbosity == "verbose"
+
+    def next_session_id(self) -> str:
+        return f"session-{next(self._session_counter)}"
+
+    async def emit(
+        self,
+        event: str,
+        *,
+        session: InstrumentedServerSession | Any | None = None,
+        request_id: str | int | None = None,
+        reason: str | None = None,
+        exception: BaseException | None = None,
+        error: str | None = None,
+        fields: dict[str, Any] | None = None,
+        verbose_fields: dict[str, Any] | None = None,
+        mirror_to_client: bool = True,
+    ) -> None:
+        if not self.enabled:
+            return
+
+        payload: dict[str, Any] = {"event": event}
+        if session is not None:
+            payload.update(_session_log_fields(session))
+        if request_id is not None:
+            payload["request_id"] = request_id
+        if reason:
+            payload["reason"] = reason
+        if exception is not None:
+            payload["exception_type"] = type(exception).__name__
+        if error is None and exception is not None:
+            error = str(exception)
+        if error:
+            payload["error"] = _clean_text(error, limit=200)
+        payload.update(fields or {})
+        if self.verbose and verbose_fields:
+            payload.update(verbose_fields)
+
+        message = _format_log_line(payload)
+        self.logger.info(message)
+
+        if not mirror_to_client or session is None or not getattr(session, "can_mirror_lifecycle_logs", False):
+            return
+        try:
+            await session.send_log_message(
+                level="info",
+                data=message,
+                logger="codepreproc.mcp.lifecycle",
+                related_request_id=request_id,
+            )
+        except Exception:
+            self.logger.debug("failed_to_send_lifecycle_log event=%s", event, exc_info=True)
+
+
+class InstrumentedServerSession(ServerSession):
+    """ServerSession with explicit lifecycle logging hooks."""
+
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        init_options: InitializationOptions,
+        *,
+        lifecycle_logger: MCPConnectionLifecycleLogger,
+        session_id: str,
+        stateless: bool = False,
+    ) -> None:
+        super().__init__(read_stream, write_stream, init_options, stateless=stateless)
+        self.lifecycle_logger = lifecycle_logger
+        self.session_id = session_id
+        self.started_at = time.monotonic()
+        self.initialize_started_at: float | None = None
+        self.initialize_completed_at: float | None = None
+        self.transport_close_reason: str | None = None
+        self.transport_close_error: str | None = None
+        self.transport_close_exception_type: str | None = None
+
+    @property
+    def uptime_ms(self) -> int:
+        return int((time.monotonic() - self.started_at) * 1000)
+
+    @property
+    def client_name(self) -> str | None:
+        return self._client_params.clientInfo.name if self._client_params else None
+
+    @property
+    def client_version(self) -> str | None:
+        return self._client_params.clientInfo.version if self._client_params else None
+
+    @property
+    def protocol_version(self) -> str | int | None:
+        return self._client_params.protocolVersion if self._client_params else None
+
+    @property
+    def can_mirror_lifecycle_logs(self) -> bool:
+        return self._initialization_state == InitializationState.Initialized
+
+    async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
+        match responder.request.root:
+            case types.InitializeRequest(params=params):
+                self.initialize_started_at = time.monotonic()
+                await self.lifecycle_logger.emit(
+                    "initialize_received",
+                    session=self,
+                    request_id=responder.request_id,
+                    fields={
+                        "client_name": params.clientInfo.name,
+                        "client_version": params.clientInfo.version,
+                        "protocol_version": params.protocolVersion,
+                    },
+                    verbose_fields={
+                        "client_capabilities": _capabilities_summary(params.capabilities),
+                        "roots_supported": bool(params.capabilities and params.capabilities.roots),
+                    },
+                    mirror_to_client=False,
+                )
+                try:
+                    await super()._received_request(responder)
+                except Exception as exc:
+                    await self.lifecycle_logger.emit(
+                        "initialize_failed",
+                        session=self,
+                        request_id=responder.request_id,
+                        reason="initialize_request_failed",
+                        exception=exc,
+                        fields={
+                            "client_name": params.clientInfo.name,
+                            "client_version": params.clientInfo.version,
+                            "protocol_version": params.protocolVersion,
+                        },
+                        verbose_fields={
+                            "client_capabilities": _capabilities_summary(params.capabilities),
+                        },
+                        mirror_to_client=False,
+                    )
+                    raise
+                self.initialize_completed_at = time.monotonic()
+                await self.lifecycle_logger.emit(
+                    "initialize_completed",
+                    session=self,
+                    request_id=responder.request_id,
+                    fields={
+                        "client_name": params.clientInfo.name,
+                        "client_version": params.clientInfo.version,
+                        "protocol_version": params.protocolVersion,
+                        "initialization_ms": int((self.initialize_completed_at - self.initialize_started_at) * 1000),
+                    },
+                    verbose_fields={
+                        "client_capabilities": _capabilities_summary(params.capabilities),
+                        "roots_supported": bool(params.capabilities and params.capabilities.roots),
+                    },
+                )
+            case types.PingRequest():
+                await super()._received_request(responder)
+            case _:
+                if self._initialization_state != InitializationState.Initialized:
+                    message = "Received request before initialization was complete"
+                    await self.lifecycle_logger.emit(
+                        "request_before_initialize",
+                        session=self,
+                        request_id=responder.request_id,
+                        reason="request_before_initialize",
+                        error=message,
+                        fields={"request_type": type(responder.request.root).__name__},
+                        mirror_to_client=False,
+                    )
+                    raise RuntimeError(message)
+                await super()._received_request(responder)
+
+    async def _received_notification(self, notification: types.ClientNotification) -> None:
+        await super()._received_notification(notification)
+        if isinstance(notification.root, types.InitializedNotification):
+            await self.lifecycle_logger.emit(
+                "initialized_notification_received",
+                session=self,
+                verbose_fields={
+                    "client_capabilities": _capabilities_summary(self._client_params.capabilities)
+                    if self._client_params
+                    else "unknown",
+                },
+            )
+
+    async def _receive_loop(self) -> None:
+        transport_reason = "client_eof"
+        transport_exception: BaseException | None = None
+        async with (self._read_stream, self._write_stream):
+            try:
+                async for message in self._read_stream:
+                    if isinstance(message, Exception):  # pragma: no cover
+                        await self._handle_incoming(message)
+                    elif isinstance(message.message.root, types.JSONRPCRequest):
+                        try:
+                            validated_request = self._receive_request_type.model_validate(
+                                message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True)
+                            )
+                            responder = RequestResponder(
+                                request_id=message.message.root.id,
+                                request_meta=validated_request.root.params.meta if validated_request.root.params else None,
+                                request=validated_request,
+                                session=self,
+                                on_complete=lambda r: self._in_flight.pop(r.request_id, None),
+                                message_metadata=message.metadata,
+                            )
+                            self._in_flight[responder.request_id] = responder
+                            await self._received_request(responder)
+
+                            if not responder._completed:  # type: ignore[reportPrivateUsage]
+                                await self._handle_incoming(responder)
+                        except Exception as exc:
+                            sdk_server_logger.warning("Failed to validate request: %s", exc)
+                            sdk_server_logger.debug("Message that failed validation: %s", message.message.root)
+                            error_response = types.JSONRPCError(
+                                jsonrpc="2.0",
+                                id=message.message.root.id,
+                                error=types.ErrorData(
+                                    code=types.INVALID_PARAMS,
+                                    message="Invalid request parameters",
+                                    data="",
+                                ),
+                            )
+                            session_message = SessionMessage(message=types.JSONRPCMessage(error_response))
+                            await self._write_stream.send(session_message)
+
+                    elif isinstance(message.message.root, types.JSONRPCNotification):
+                        try:
+                            notification = self._receive_notification_type.model_validate(
+                                message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True)
+                            )
+                            if isinstance(notification.root, types.CancelledNotification):
+                                cancelled_id = notification.root.params.requestId
+                                if cancelled_id in self._in_flight:
+                                    await self._in_flight[cancelled_id].cancel()
+                            else:
+                                if isinstance(notification.root, types.ProgressNotification):  # pragma: no cover
+                                    progress_token = notification.root.params.progressToken
+                                    if progress_token in self._progress_callbacks:
+                                        callback = self._progress_callbacks[progress_token]
+                                        try:
+                                            await callback(
+                                                notification.root.params.progress,
+                                                notification.root.params.total,
+                                                notification.root.params.message,
+                                            )
+                                        except Exception as exc:  # pragma: no cover
+                                            sdk_server_logger.error("Progress callback raised an exception: %s", exc)
+                                await self._received_notification(notification)
+                                await self._handle_incoming(notification)
+                        except Exception as exc:  # pragma: no cover
+                            sdk_server_logger.warning(
+                                "Failed to validate notification: %s. Message was: %s",
+                                exc,
+                                message.message.root,
+                            )
+                    else:
+                        await self._handle_response(message)
+
+            except anyio.ClosedResourceError as exc:
+                transport_reason = "read_stream_closed"
+                transport_exception = exc
+            except Exception as exc:  # pragma: no cover
+                transport_reason = "receive_loop_exception"
+                transport_exception = exc
+                await self.lifecycle_logger.emit(
+                    "unhandled_receive_exception",
+                    session=self,
+                    reason=transport_reason,
+                    exception=exc,
+                    mirror_to_client=False,
+                )
+            finally:
+                self.transport_close_reason = transport_reason
+                if transport_exception is not None:
+                    self.transport_close_exception_type = type(transport_exception).__name__
+                    self.transport_close_error = _clean_text(str(transport_exception), limit=200)
+                pending_response_count = len(self._response_streams)
+                in_flight_request_count = len(self._in_flight)
+
+                for response_id, stream in list(self._response_streams.items()):
+                    error = types.ErrorData(code=types.CONNECTION_CLOSED, message="Connection closed")
+                    try:
+                        await stream.send(types.JSONRPCError(jsonrpc="2.0", id=response_id, error=error))
+                        await stream.aclose()
+                    except Exception:  # pragma: no cover
+                        pass
+                self._response_streams.clear()
+
+                await self.lifecycle_logger.emit(
+                    "transport_closed",
+                    session=self,
+                    reason=transport_reason,
+                    exception=transport_exception,
+                    verbose_fields={
+                        "pending_response_count": pending_response_count,
+                        "in_flight_request_count": in_flight_request_count,
+                        "transport_reason_detail": transport_reason,
+                    },
+                )
+
+
+class InstrumentedServer(Server):
+    """Server wrapper that instantiates InstrumentedServerSession and logs lifecycle events."""
+
+    def __init__(self, *args: Any, lifecycle_logger: MCPConnectionLifecycleLogger, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.lifecycle_logger = lifecycle_logger
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        initialization_options: InitializationOptions,
+        raise_exceptions: bool = False,
+        stateless: bool = False,
+    ) -> None:
+        async with AsyncExitStack() as stack:
+            lifespan_context = await stack.enter_async_context(self.lifespan(self))
+            session = await stack.enter_async_context(
+                InstrumentedServerSession(
+                    read_stream,
+                    write_stream,
+                    initialization_options,
+                    stateless=stateless,
+                    lifecycle_logger=self.lifecycle_logger,
+                    session_id=self.lifecycle_logger.next_session_id(),
+                )
+            )
+            await self.lifecycle_logger.emit(
+                "transport_open",
+                session=session,
+                verbose_fields={"stateless": stateless},
+                mirror_to_client=False,
+            )
+
+            task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
+            if task_support is not None:
+                task_support.configure_session(session)
+                await stack.enter_async_context(task_support.run())
+
+            async with anyio.create_task_group() as tg:
+                try:
+                    async for message in session.incoming_messages:
+                        sdk_server_logger.debug("Received message: %s", message)
+                        tg.start_soon(
+                            self._handle_message,
+                            message,
+                            session,
+                            lifespan_context,
+                            raise_exceptions,
+                        )
+                finally:
+                    tg.cancel_scope.cancel()
+                    await self.lifecycle_logger.emit(
+                        "session_shutdown",
+                        session=session,
+                        reason=session.transport_close_reason or "server_run_complete",
+                        verbose_fields={
+                            "in_flight_request_count": len(session._in_flight),
+                            "pending_response_count": len(session._response_streams),
+                        },
+                    )
+
+    async def _handle_request(
+        self,
+        message: RequestResponder[types.ClientRequest, types.ServerResult],
+        req: types.ClientRequestType,
+        session: ServerSession,
+        lifespan_context: Any,
+        raise_exceptions: bool,
+    ):
+        sdk_server_logger.info("Processing request of type %s", type(req).__name__)
+
+        if handler := self.request_handlers.get(type(req)):
+            sdk_server_logger.debug("Dispatching request of type %s", type(req).__name__)
+
+            token = None
+            try:
+                request_data = None
+                close_sse_stream_cb = None
+                close_standalone_sse_stream_cb = None
+                if message.message_metadata is not None and isinstance(message.message_metadata, ServerMessageMetadata):
+                    request_data = message.message_metadata.request_context
+                    close_sse_stream_cb = message.message_metadata.close_sse_stream
+                    close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
+
+                client_capabilities = session.client_params.capabilities if session.client_params else None
+                task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
+                task_metadata = None
+                if hasattr(req, "params") and req.params is not None:
+                    task_metadata = getattr(req.params, "task", None)
+                token = request_ctx.set(
+                    RequestContext(
+                        message.request_id,
+                        message.request_meta,
+                        session,
+                        lifespan_context,
+                        Experimental(
+                            task_metadata=task_metadata,
+                            _client_capabilities=client_capabilities,
+                            _session=session,
+                            _task_support=task_support,
+                        ),
+                        request=request_data,
+                        close_sse_stream=close_sse_stream_cb,
+                        close_standalone_sse_stream=close_standalone_sse_stream_cb,
+                    )
+                )
+                response = await handler(req)
+            except McpError as err:  # pragma: no cover
+                response = err.error
+            except anyio.get_cancelled_exc_class():
+                if message.cancelled:
+                    sdk_server_logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
+                    return
+                raise
+            except Exception as err:  # pragma: no cover
+                if raise_exceptions:
+                    raise err
+                response = types.ErrorData(code=0, message=str(err), data=None)
+            finally:
+                if token is not None:
+                    request_ctx.reset(token)
+        else:  # pragma: no cover
+            response = types.ErrorData(
+                code=types.METHOD_NOT_FOUND,
+                message="Method not found",
+            )
+
+        try:
+            await message.respond(response)
+        except (anyio.BrokenResourceError, anyio.ClosedResourceError) as exc:
+            await self.lifecycle_logger.emit(
+                "response_dropped_transport_closed",
+                session=session,
+                request_id=message.request_id,
+                reason="transport_closed_before_response",
+                exception=exc,
+            )
+            sdk_server_logger.debug("Response for %s dropped - transport closed", message.request_id)
+            return
+
+        sdk_server_logger.debug("Response sent")
+
+
+async def resolve_session_roots(
+    *,
+    session: ServerSession,
+    lifecycle_logger: MCPConnectionLifecycleLogger,
+    request_id: str | int | None,
+) -> list[str]:
+    """Resolve client roots and log failures without breaking tool calls."""
+
+    if not (session.client_params and session.client_params.capabilities and session.client_params.capabilities.roots):
+        return []
+    try:
+        roots_result = await session.list_roots()
+    except Exception as exc:
+        await lifecycle_logger.emit(
+            "roots_resolution_failed",
+            session=session,
+            request_id=request_id,
+            reason="list_roots_failed",
+            exception=exc,
+            verbose_fields={"roots_supported": True},
+        )
+        return []
+    return [str(root.uri) for root in roots_result.roots if str(root.uri).startswith("file://")]
+
+
+def _session_log_fields(session: InstrumentedServerSession | Any) -> dict[str, Any]:
+    payload: dict[str, Any] = {}
+    session_id = getattr(session, "session_id", None)
+    if session_id is not None:
+        payload["session_id"] = session_id
+    client_name = getattr(session, "client_name", None)
+    if client_name:
+        payload["client_name"] = client_name
+    client_version = getattr(session, "client_version", None)
+    if client_version:
+        payload["client_version"] = client_version
+    protocol_version = getattr(session, "protocol_version", None)
+    if protocol_version is not None:
+        payload["protocol_version"] = protocol_version
+    uptime_ms = getattr(session, "uptime_ms", None)
+    if uptime_ms is not None:
+        payload["uptime_ms"] = uptime_ms
+    return payload
+
+
+def _capabilities_summary(capabilities: types.ClientCapabilities | None) -> str:
+    if capabilities is None:
+        return "none"
+    names: list[str] = []
+    if capabilities.roots is not None:
+        names.append("roots")
+    if capabilities.sampling is not None:
+        names.append("sampling")
+    if capabilities.elicitation is not None:
+        names.append("elicitation")
+    if capabilities.experimental:
+        names.append("experimental")
+    if capabilities.tasks is not None:
+        names.append("tasks")
+    return ",".join(names) or "none"
+
+
+def _clean_text(value: str, *, limit: int) -> str:
+    return " ".join(value.split())[:limit]
+
+
+def _format_log_line(payload: dict[str, Any]) -> str:
+    parts = ["mcp_lifecycle"]
+    for key, value in payload.items():
+        if value is None:
+            continue
+        parts.append(f"{key}={_format_value(value)}")
+    return " ".join(parts)
+
+
+def _format_value(value: Any) -> str:
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, (int, float)):
+        return str(value)
+    if isinstance(value, (dict, list)):
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
+    text = str(value)
+    if not text or any(char.isspace() for char in text) or any(char in {'"', "="} for char in text):
+        return json.dumps(text, ensure_ascii=False)
+    return text