aiohttp-tiny-mcp 0.1.0__py3-none-any.whl

aiohttp_tiny_mcp/exchange.py

@@ -0,0 +1,279 @@
+"""Per-request context: dependency resolution, progress, MRTR access."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import logging
+from collections.abc import Awaitable, Callable, Mapping
+from contextlib import AsyncExitStack, asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from aiohttp import web
+
+from .adapter import Adapter, RequestLike
+from .auth import Principal
+from .core import (
+    Answer,
+    AnswerAction,
+    Call,
+    ClientInfo,
+    InputRequest,
+    NeedInput,
+    answer_of,
+    logs_at,
+)
+from .hub import ASK, Hub, topic
+from .sessions import Session, SessionAccess
+from .sessions import new_session_id as new_id
+from .sse import SSEResponse
+
+if TYPE_CHECKING:
+    from .registry import Registry
+
+log = logging.getLogger("aiohttp_tiny_mcp")
+
+
+def is_reply(body: Any) -> bool:
+    """Pushed-question answers are JSON-RPC responses without a method."""
+    return (
+        isinstance(body, dict)
+        and "method" not in body
+        and ("result" in body or "error" in body)
+        and isinstance(body.get("id"), str)
+    )
+
+
+async def relay_reply(hub: Hub, body: Mapping[str, Any]) -> None:
+    """Publish an answer to its namespaced question topic, which may be read on another node."""
+    reply = body["result"] if "result" in body else body["error"]
+    await hub.publish(topic(ASK, str(body["id"])), {"reply": reply})
+
+
+@dataclass(frozen=True, slots=True)
+class Instance:
+    """An existing dependency object registered with Registry.provide_instance."""
+
+    value: Any
+
+
+class Exchange:
+    def __init__(
+        self,
+        registry: Registry,
+        request: RequestLike,
+        adapter: Adapter,
+        call: Call,
+        session: Session | None = None,
+    ) -> None:
+        self.registry = registry
+        self.request = request
+        self.adapter = adapter
+        self.call = call
+        #: Application values kept between calls, or None where this request
+        #: reached no session.
+        self.session = session
+        self.principal: Principal | None = None
+        self.keep_log_level: Callable[[str], None] | None = None
+        self.sse: SSEResponse | None = None
+        self.send: Callable[[Mapping[str, Any]], Awaitable[None]] | None = None
+        self.stack = AsyncExitStack()
+        self.resolved: dict[type, Any] = {}
+        self.cancelled = asyncio.Event()
+
+    @asynccontextmanager
+    async def scope(self):
+        """Release per-call dependencies, propagating handler exceptions into providers for
+        rollback.
+        """
+        async with self.stack:
+            yield
+
+    async def resolve(self, kind: type) -> Any:
+        if kind is Exchange:
+            return self
+        if kind is Principal:
+            return self.principal
+        if kind in self.resolved:
+            return self.resolved[kind]
+        source = self.registry.providers.get(kind)
+        if source is None:
+            raise LookupError(f"no provider for {kind!r}")
+        if isinstance(source, Instance):
+            value = source.value
+        elif isinstance(source, web.AppKey):
+            value = self.request.app[source]
+        elif inspect.isasyncgenfunction(source):
+            value = await self.stack.enter_async_context(asynccontextmanager(source)(self))
+        else:
+            value = await source(self)
+        self.resolved[kind] = value
+        return value
+
+    @property
+    def id(self) -> str | int:
+        assert self.call.id is not None
+        return self.call.id
+
+    @property
+    def can_ask(self) -> bool:
+        """Whether input can be requested: MRTR requires declared elicitation; the tool-argument
+        fallback does not.
+        """
+        if self.adapter.asks_in_arguments:
+            return True
+        return self.asks_somehow and self.declared_elicitation
+
+    @property
+    def asks_somehow(self) -> bool:
+        """Whether this revision can put a question to a client at all."""
+        return self.adapter.can_ask or self.adapter.can_push_ask or self.adapter.asks_in_arguments
+
+    @property
+    def declared_elicitation(self) -> bool:
+        """Capability presence is sufficient; an empty mapping is a valid declaration."""
+        return self.call.client.capabilities.get("elicitation") is not None
+
+    @property
+    def sessions(self) -> SessionAccess:
+        """Access sessions by explicit handles, including on revisions without protocol sessions."""
+        return SessionAccess(self.registry.session_store, self.registry.session_ttl_seconds)
+
+    @property
+    def answers(self) -> dict[str, Any]:
+        return dict(self.call.answers)
+
+    def answered(self, key: str) -> bool:
+        """Check for a reply, including accepted forms with empty content."""
+        return key in self.call.actions or key in self.call.answers
+
+    def accepted(self, key: str) -> bool:
+        """Check the action; accepted and declined replies can both have empty content."""
+        return self.call.actions.get(key) is AnswerAction.ACCEPT
+
+    async def ask(
+        self,
+        key: str,
+        request: InputRequest,
+        *,
+        default: Mapping[str, Any] | None = None,
+    ) -> Answer:
+        """Ask for input and return the answer.
+
+        MRTR raises NeedInput and restarts the handler when the answer arrives, possibly on
+        another node. Put irreversible work after the final ask; preceding work may run again.
+
+        For clients that cannot be asked, `default` accepts an elicit_accept/decline/cancel
+        result. Without a default, the call fails.
+        """
+        if self.answered(key):
+            return Answer(
+                action=self.call.actions.get(key, AnswerAction.ACCEPT),
+                content=self.call.answers.get(key) or {},
+            )
+        if self.can_ask and self.adapter.can_push_ask:
+            return await self.push_ask(key, request)
+        if default is not None and not self.can_ask:
+            return answer_of(default)
+        raise NeedInput({key: request})
+
+    async def push_ask(self, key: str, request: InputRequest) -> Answer:
+        """Send a question on the active stream and wait through the hub.
+
+        The reply may reach another node. Capture the cursor before sending so fast replies are
+        not missed.
+        """
+        hub = self.registry.hub
+        wire_id = new_id()
+        where = topic(ASK, wire_id)
+        cursor = await hub.position(where)
+        await self.emit({"jsonrpc": "2.0", "id": wire_id, **request})
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + self.registry.ask_timeout_seconds
+        try:
+            while not self.cancelled.is_set():
+                left = deadline - loop.time()
+                if left <= 0:
+                    return Answer(action=AnswerAction.CANCEL)
+                messages, cursor = await hub.poll(
+                    where, cursor, timeout=min(left, self.registry.hub_poll_seconds)
+                )
+                for message in messages:
+                    reply = message.get("reply")
+                    if isinstance(reply, Mapping):
+                        return answer_of(reply)
+            return Answer(action=AnswerAction.CANCEL)
+        finally:
+            await hub.delete(where)
+
+    def action(self, key: str) -> AnswerAction | None:
+        """Return ACCEPT, DECLINE, CANCEL, or None if unanswered."""
+        return self.call.actions.get(key)
+
+    @property
+    def state(self) -> Any:
+        return self.call.state
+
+    @property
+    def client_info(self) -> ClientInfo:
+        return self.call.client
+
+    @property
+    def log_level(self) -> str | None:
+        return self.call.log_level
+
+    @property
+    def progress_token(self) -> str | int:
+        return self.call.progress_token if self.call.progress_token is not None else self.id
+
+    async def emit(self, payload: Mapping[str, Any]) -> None:
+        if self.cancelled.is_set():
+            return
+        if self.send is not None:
+            await self.send(payload)
+            return
+        assert self.sse is not None
+        text = json.dumps(payload, ensure_ascii=False)
+        log.debug("-> [%s] %s", self.adapter.version, text)
+        await self.sse.send(text)
+
+    def cancel(self) -> None:
+        self.cancelled.set()
+
+    async def wait_cancelled(self) -> None:
+        await self.cancelled.wait()
+
+    def logs(self, level: str) -> bool:
+        """Whether a message of `level` would reach this client."""
+        return logs_at(self.log_level, level)
+
+    async def log(self, level: str, data: Any, *, logger: str | None = None) -> None:
+        """Emit a message at the requested severity on the current request stream. Without a
+        stream, emit nothing.
+        """
+        if not self.logs(level):
+            return
+        params: dict[str, Any] = {"level": level, "data": data}
+        if logger is not None:
+            params["logger"] = logger
+        await self.emit({"jsonrpc": "2.0", "method": "notifications/message", "params": params})
+
+    async def progress(
+        self, progress: float, total: float | None = None, message: str | None = None
+    ) -> None:
+        """Emit on the current HTTP response or stdio request."""
+        if self.sse is None and self.send is None:
+            return
+        params: dict[str, Any] = {"progressToken": self.progress_token, "progress": progress}
+        if total is not None:
+            params["total"] = total
+        if message is not None and self.adapter.progress_message:
+            params["message"] = message
+        await self.emit({"jsonrpc": "2.0", "method": "notifications/progress", "params": params})
+
+    def open(self, *, compress: bool = True, **headers: str) -> SSEResponse:
+        """The response this call streams on. Prepared by the caller."""
+        self.sse = SSEResponse(compress=compress, headers=headers)
+        return self.sse

aiohttp_tiny_mcp/http_sse.py

@@ -0,0 +1,267 @@
+"""The legacy 2024-11-05 HTTP+SSE transport."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from contextlib import suppress
+from dataclasses import replace
+
+from aiohttp import web
+
+from .adapter import Adapter
+from .core import (
+    Call,
+    DecodeFailure,
+    Failure,
+    FailureKind,
+    Operation,
+    Preamble,
+    Rejected,
+    Value,
+)
+from .dispatcher import Dispatcher
+from .endpoint import Endpoint
+from .exchange import Exchange
+from .hub import topic
+from .namespaces import scoped
+from .protocol.selection import AdapterSet
+from .registry import Registry
+from .sessions import (
+    Session,
+    SessionRecord,
+    handshake_data,
+    new_session_id,
+    stored_capabilities,
+    stored_log_level,
+    stored_version,
+)
+from .sse import SSEResponse
+
+log = logging.getLogger(__name__)
+
+STREAM = "sse"
+
+VERSION = "2024-11-05"
+
+
+class SseEndpoint:
+    """Two endpoints: one to listen on, one to send to."""
+
+    def __init__(
+        self,
+        registry: Registry,
+        *,
+        adapters: AdapterSet | None = None,
+        allowed_origins: set[str] | None = None,
+        trust_proxy_origin_validation: bool = False,
+        compress: bool = True,
+        sse_path: str = "/sse",
+        message_path: str = "/messages",
+    ) -> None:
+        self.registry = registry
+        self.adapters = adapters or AdapterSet.default()
+        self.dispatcher = Dispatcher(registry)
+        self.origins = Endpoint(
+            registry,
+            adapters=self.adapters,
+            allowed_origins=allowed_origins,
+            trust_proxy_origin_validation=trust_proxy_origin_validation,
+        )
+        self.compress = compress
+        self.sse_path = sse_path
+        self.message_path = message_path
+
+    @property
+    def adapter(self) -> Adapter:
+        """The one that renders a failure before a revision is known."""
+        return self.adapters.by_version[VERSION]
+
+    def speaking(self, pre: Preamble, record: SessionRecord) -> Adapter:
+        """Select a handshake-capable revision."""
+        chosen = self.adapters.select(pre, stored_version(record) or VERSION)
+        if not chosen.has_handshake:
+            raise Rejected(
+                Failure(
+                    FailureKind.UNSUPPORTED_VERSION,
+                    f"{chosen.version} is not spoken over this transport",
+                    data={
+                        "supported": [a.version for a in self.adapters.adapters if a.has_handshake]
+                    },
+                )
+            )
+        return chosen
+
+    async def listen(self, request: web.Request) -> web.StreamResponse:
+        """Open a stream and send its POST address."""
+        try:
+            self.origins.check_origin(request)
+        except Rejected as e:
+            return self.origins.render_failure(self.adapter, e.failure)
+
+        session_id = await self.open_session()
+        where = topic(STREAM, session_id)
+        hub = self.registry.hub
+        cursor = await hub.position(where)
+
+        response = SSEResponse(compress=self.compress)
+        await response.prepare(request)
+        posting = self.mounted(request) + self.message_path
+        await self.write(response, "endpoint", f"{posting}?session_id={session_id}")
+
+        try:
+            await self.relay_until_disconnect(request, response, where, cursor)
+        finally:
+            await hub.delete(where)
+            await self.registry.session_store.delete(scoped(session_id))
+        return response
+
+    async def open_session(self) -> str:
+        """Create the short-lived session represented by an open stream."""
+        session_id = new_session_id()
+        await self.registry.session_store.create(
+            scoped(session_id),
+            handshake_data(VERSION, {}),
+            ttl_seconds=self.registry.session_ttl_seconds,
+        )
+        return session_id
+
+    async def relay_until_disconnect(
+        self, request: web.Request, response: SSEResponse, where: str, cursor: str
+    ) -> None:
+        """Run a relay until its client disconnects, then always join it."""
+        relay = asyncio.create_task(self.relay(response, where, cursor))
+        try:
+            while not relay.done():
+                await asyncio.wait({relay}, timeout=0.05)
+                transport = request.transport
+                if transport is None or transport.is_closing():
+                    break
+        finally:
+            relay.cancel()
+            with suppress(asyncio.CancelledError, ConnectionError):
+                await relay
+
+    async def relay(self, response: SSEResponse, where: str, cursor: str) -> None:
+        """Write everything published for this connection, until cancelled."""
+        hub = self.registry.hub
+        while True:
+            messages, cursor = await hub.poll(where, cursor, timeout=self.registry.hub_poll_seconds)
+            for payload in messages:
+                await self.write(response, "message", json.dumps(payload, ensure_ascii=False))
+
+    async def write(self, response: SSEResponse, event: str, data: str) -> None:
+        log.debug("-> [%s] %s %s", VERSION, event, data)
+        await response.send(data, event=event)
+
+    async def receive(self, request: web.Request) -> web.StreamResponse:
+        """Take one message and answer 202. The reply goes to the stream."""
+        try:
+            self.origins.check_origin(request)
+        except Rejected as e:
+            return self.origins.render_failure(self.adapter, e.failure)
+
+        session_id = request.query.get("session_id", "")
+        record = await self.registry.session_store.get(scoped(session_id)) if session_id else None
+        if record is None:
+            return web.json_response({"error": "no such session"}, status=404)
+
+        raw = await request.read()
+        log.debug("<- [%s] %s", VERSION, raw.decode("utf-8", "replace"))
+        where = topic(STREAM, session_id)
+
+        pre = Preamble.of(raw, request.headers)
+        try:
+            adapter = self.speaking(pre, record)
+            items = adapter.decode(pre)
+        except Rejected as e:
+            await self.registry.hub.publish(where, self.adapter.encode_failure(None, e.failure))
+            return web.Response(status=202)
+
+        for item in items:
+            await self.serve(adapter, item, session_id, record, where)
+        return web.Response(status=202)
+
+    async def serve(
+        self,
+        adapter: Adapter,
+        item: Call | DecodeFailure,
+        session_id: str,
+        record: SessionRecord,
+        where: str,
+    ) -> None:
+        hub = self.registry.hub
+        if isinstance(item, DecodeFailure):
+            if item.must_respond:
+                await hub.publish(where, adapter.encode_failure(item.id, item.failure))
+            return
+
+        item.client = replace(
+            item.client,
+            capabilities={**stored_capabilities(record), **item.client.capabilities},
+        )
+        if item.log_level is None:
+            item.log_level = stored_log_level(record)
+
+        session = Session(
+            self.registry.session_store, session_id, record, self.registry.session_ttl_seconds
+        )
+        ex = Exchange(self.registry, self.origins, adapter, item, session=session)
+        ex.send = lambda payload: hub.publish(where, payload)
+
+        try:
+            outcome = await self.dispatcher.run(ex)
+        except Exception as e:  # noqa: BLE001 -- a failed call must still answer
+            log.exception("sse request failed")
+            outcome = Failure(FailureKind.INTERNAL, f"{type(e).__name__}: {e}")
+
+        if item.operation is Operation.DESCRIBE and isinstance(outcome, Value):
+            negotiated = getattr(outcome.result, "protocol_version", None)
+            if isinstance(negotiated, str):
+                await session.remember_version(negotiated)
+
+        if item.is_notification:
+            return
+        await hub.publish(where, adapter.encode(item, self.registry, outcome))
+
+    def mounted(self, request: web.Request) -> str:
+        """The prefix the client reached this stream through.
+
+        A subapplication adds its prefix to the request path and not to the
+        route, and the address on the stream is one a client posts to, so it
+        has to carry that prefix.
+        """
+        path = request.path
+        return path[: -len(self.sse_path)] if path.endswith(self.sse_path) else ""
+
+    def routes(
+        self, sse_path: str | None = None, message_path: str | None = None
+    ) -> list[web.RouteDef]:
+        """One route to listen on, one to post to.
+
+        Either path may be set here or on the constructor. Both are kept,
+        because the stream names the posting path to the client.
+        """
+        if sse_path is not None:
+            self.sse_path = sse_path
+        if message_path is not None:
+            self.message_path = message_path
+        log.debug("HTTP+SSE stream at %s, messages at %s", self.sse_path, self.message_path)
+        return [
+            web.get(self.sse_path, self.listen),
+            web.post(self.message_path, self.receive),
+        ]
+
+    def setup(
+        self,
+        app: web.Application,
+        sse_path: str | None = None,
+        message_path: str | None = None,
+    ) -> web.Application:
+        log.debug("adding the HTTP+SSE routes to %r", app)
+        app.add_routes(self.routes(sse_path, message_path))
+        return app
+
+
+__all__ = ["STREAM", "VERSION", "SseEndpoint"]

aiohttp_tiny_mcp/hub.py

@@ -0,0 +1,109 @@
+"""Cursor-based event storage shared by reply and notification streams.
+
+Applications supply the backend; MemoryHub serves one process. Shared storage lets one node
+publish replies that another is awaiting.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from abc import abstractmethod
+from collections.abc import Mapping, Sequence
+from typing import Any, Protocol, runtime_checkable
+
+from .namespaces import scoped
+
+Cursor = str
+
+START: Cursor = ""
+
+NOTIFICATIONS = "notifications"
+
+ASK = "ask"
+
+
+@runtime_checkable
+class Hub(Protocol):
+    """Application-provided event storage.
+
+    Subclass it to have the methods checked and the missing ones refused, or
+    supply any object with these four methods: this is a protocol, so a
+    backend that inherits nothing is still a Hub.
+    """
+
+    @abstractmethod
+    async def publish(self, topic: str, message: Mapping[str, Any]) -> None:
+        """Append one message to `topic`."""
+
+    @abstractmethod
+    async def position(self, topic: str) -> Cursor:
+        """Capture the cursor before triggering a publish so replies preceding the first poll are
+        included.
+        """
+
+    @abstractmethod
+    async def poll(
+        self, topic: str, cursor: Cursor, *, timeout: float
+    ) -> tuple[Sequence[Mapping[str, Any]], Cursor]:
+        """Return messages after cursor and the next cursor.
+
+        Wait up to timeout seconds for a message. Backends choose whether to poll or wait for
+        notification; timeout is a deadline, not a polling interval.
+        """
+
+    @abstractmethod
+    async def delete(self, topic: str) -> None:
+        """Delete a completed topic."""
+
+
+def topic(kind: str, name: str | None = None) -> str:
+    """Prefix the topic with the current namespace to isolate callers."""
+    return scoped(kind if name is None else f"{kind}/{name}")
+
+
+class MemoryHub(Hub):
+    """Single-process hub, waiting on a condition rather than sleeping."""
+
+    def __init__(self) -> None:
+        self.rows: dict[str, list[tuple[int, Mapping[str, Any]]]] = {}
+        self.last_id = 0
+        self.arrived = asyncio.Condition()
+
+    async def publish(self, topic: str, message: Mapping[str, Any]) -> None:
+        async with self.arrived:
+            self.last_id += 1
+            self.rows.setdefault(topic, []).append((self.last_id, dict(message)))
+            self.arrived.notify_all()
+
+    async def position(self, topic: str) -> Cursor:
+        rows = self.rows.get(topic)
+        return str(rows[-1][0]) if rows else START
+
+    def after(self, topic: str, cursor: Cursor) -> list[tuple[int, Mapping[str, Any]]]:
+        least = int(cursor) if cursor else 0
+        return [row for row in self.rows.get(topic, []) if row[0] > least]
+
+    async def poll(
+        self, topic: str, cursor: Cursor, *, timeout: float
+    ) -> tuple[Sequence[Mapping[str, Any]], Cursor]:
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + timeout
+        async with self.arrived:
+            while True:
+                found = self.after(topic, cursor)
+                if found:
+                    return [message for _, message in found], str(found[-1][0])
+                left = deadline - loop.time()
+                if left <= 0:
+                    return [], cursor
+                try:
+                    await asyncio.wait_for(self.arrived.wait(), left)
+                except asyncio.TimeoutError:
+                    return [], cursor
+
+    async def delete(self, topic: str) -> None:
+        async with self.arrived:
+            self.rows.pop(topic, None)
+
+
+__all__ = ["ASK", "NOTIFICATIONS", "START", "Cursor", "Hub", "MemoryHub", "topic"]