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

aiohttp_tiny_mcp/registry.py

@@ -0,0 +1,162 @@
+"""Public API: declare tools, resources and prompts once, revision-free."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from .auth import Authorization, Principal
+from .exchange import Exchange, Instance
+from .hub import Hub, MemoryHub
+from .models import Implementation
+from .request_state import DEFAULT_TTL_SECONDS as STATE_TTL_SECONDS
+from .request_state import RequestStates
+from .sessions import DEFAULT_TTL_SECONDS, MemorySessionStore, SessionStore
+from .specs import Bound, PromptSpec, ResourceSpec, ToolSpec
+
+
+class Registry:
+    """Register handlers as decorators or by passing functions directly.
+
+        @registry.tool
+        async def add(args: Add) -> int: ...
+
+        registry.tool(add, name="sum")
+        registry.resource("config://app", config)
+
+    Uses process-local memory backends by default. Supply shared `hub` and
+    `session_store` implementations before running more than one worker.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        version: str,
+        *,
+        hub: Hub | None = None,
+        session_store: SessionStore | None = None,
+        auth: Authorization | None = None,
+        instructions: str | None = None,
+        session_ttl_seconds: int = DEFAULT_TTL_SECONDS,
+        request_state_ttl_seconds: int = STATE_TTL_SECONDS,
+        hub_poll_seconds: float = 30.0,
+        ask_timeout_seconds: float = 120.0,
+    ) -> None:
+        self.info = Implementation(name=name, version=version)
+        self.instructions = instructions
+        self.hub = hub if hub is not None else MemoryHub()
+        self.session_store = session_store if session_store is not None else MemorySessionStore()
+        self.auth = auth
+        self.session_ttl_seconds = session_ttl_seconds
+        self.request_state = RequestStates(self.session_store, request_state_ttl_seconds)
+        self.hub_poll_seconds = hub_poll_seconds
+        self.ask_timeout_seconds = ask_timeout_seconds
+        self.tools: dict[str, ToolSpec] = {}
+        self.resources_fixed: dict[str, ResourceSpec] = {}
+        self.resources_templated: list[ResourceSpec] = []
+        self.prompts: dict[str, PromptSpec] = {}
+        self.completer: Bound | None = None
+        self.providers: dict[type, Any] = {}
+
+    def provide(self, kind: type, source: Any) -> None:
+        """Bind a type to where it comes from. Call before registering
+        handlers -- plans are checked at registration, not at first call.
+
+        `source` is an async callable taking the exchange, an async generator
+        for anything that must be released afterwards, or a `web.AppKey`. For
+        an object that already exists, use `provide_instance` instead.
+        """
+        self.providers[kind] = source
+
+    def provide_instance(self, value: Any, kind: type | None = None) -> None:
+        """Register an existing object under its type, or under an explicit `kind` such as a base
+        class or protocol. Use `provide` for per-request construction or cleanup.
+        """
+        if kind is None and isinstance(value, type):
+            raise TypeError(
+                f"provide_instance was given the class {value.__name__} rather than an "
+                f"instance of it. Pass the object, or say kind=type to provide the class "
+                f"itself."
+            )
+        self.providers[kind or type(value)] = Instance(value)
+
+    def check(self, what: str, bound: Bound) -> None:
+        for name, kind in bound.plan:
+            if kind is Principal:
+                if self.auth is None:
+                    raise TypeError(
+                        f"{what}: parameter {name!r} wants a Principal, but this registry "
+                        f"has no auth= -- pass an Authorization to verify tokens first"
+                    )
+                continue
+            if kind not in self.providers and kind is not Exchange:
+                label = getattr(kind, "__name__", kind)
+                raise TypeError(
+                    f"{what}: parameter {name!r} wants {label}, which has no provider -- "
+                    f"call registry.provide_instance() or registry.provide() first"
+                )
+
+    def tool(self, fn: Callable[..., Awaitable[Any]] | None = None, **kw: Any):
+        def register(fn: Callable[..., Awaitable[Any]]):
+            spec = ToolSpec.build(fn, **kw)
+            if spec.scopes and self.auth is None:
+                raise TypeError(
+                    f"{spec.name}: scopes were declared, but this registry has no auth= -- "
+                    f"a scope nothing verifies is a check that never runs"
+                )
+            self.check(spec.name, spec.bound)
+            if spec.name in self.tools:
+                raise ValueError(f"duplicate tool: {spec.name}")
+            self.tools[spec.name] = spec
+            return fn
+
+        return register(fn) if fn else register
+
+    def resource(self, uri: str, fn: Callable[..., Awaitable[Any]] | None = None, **kw: Any):
+        def register(fn: Callable[..., Awaitable[Any]]):
+            spec = ResourceSpec.build(uri, fn, **kw)
+            self.check(uri, spec.bound)
+            if spec.definition is not None:
+                if spec.definition.uri in self.resources_fixed:
+                    raise ValueError(f"duplicate resource: {spec.definition.uri}")
+                self.resources_fixed[spec.definition.uri] = spec
+            else:
+                template = spec.template.uri_template if spec.template is not None else uri
+                if any(
+                    item.template is not None and item.template.uri_template == template
+                    for item in self.resources_templated
+                ):
+                    raise ValueError(f"duplicate resource template: {template}")
+                self.resources_templated.append(spec)
+            return fn
+
+        return register(fn) if fn else register
+
+    def prompt(self, fn: Callable[..., Awaitable[Any]] | None = None, **kw: Any):
+        def register(fn: Callable[..., Awaitable[Any]]):
+            spec = PromptSpec.build(fn, **kw)
+            self.check(spec.name, spec.bound)
+            if spec.name in self.prompts:
+                raise ValueError(f"duplicate prompt: {spec.name}")
+            self.prompts[spec.name] = spec
+            return fn
+
+        return register(fn) if fn else register
+
+    def completions(self, fn: Callable[..., Awaitable[Any]]):
+        """(CompleteParams-shaped model, deps) -> list[str] | Completion."""
+        if self.completer is not None:
+            raise ValueError("duplicate completions handler")
+        self.completer = Bound.of(fn)
+        self.check("completions", self.completer)
+        return fn
+
+    def match_resource(self, uri: str) -> tuple[ResourceSpec, dict[str, str]] | None:
+        fixed = self.resources_fixed.get(uri)
+        if fixed is not None:
+            return fixed, {}
+        for spec in self.resources_templated:
+            m = spec.pattern.match(uri)
+            if m:
+                return spec, m.groupdict()
+        return None

aiohttp_tiny_mcp/request_state.py

@@ -0,0 +1,107 @@
+"""Store handler state between input rounds; clients carry only an opaque id.
+
+Ids are bound to the original call, scoped by namespace, and expire independently of sessions.
+Shared storage allows retries on another node.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import secrets
+from dataclasses import dataclass
+from typing import Any
+
+from .core import Call, Failure, FailureKind, NeedsInput
+from .namespaces import scoped
+from .sessions import SessionStore
+
+PREFIX = "state/"
+
+DEFAULT_TTL_SECONDS = 600
+
+BINDING_KEY = "binding"
+PAYLOAD_KEY = "payload"
+
+
+def binding_of(call: Call) -> str:
+    """Hash normalized call arguments and target so state cannot be reused for a different call."""
+    material = json.dumps(
+        {
+            "operation": call.operation.value,
+            "target": call.target,
+            "arguments": call.arguments,
+        },
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+    )
+    return hashlib.sha256(material.encode()).hexdigest()
+
+
+@dataclass
+class RequestStates:
+    """State rows, in the store the deployment already supplies."""
+
+    store: SessionStore
+    ttl_seconds: int = DEFAULT_TTL_SECONDS
+
+    def key(self, state_id: str) -> str:
+        return scoped(PREFIX + state_id)
+
+    async def open(self, call: Call, payload: Any) -> str:
+        """Store payload and return its opaque client-visible id."""
+        state_id = secrets.token_urlsafe(24)
+        created = await self.store.create(
+            self.key(state_id),
+            {PAYLOAD_KEY: payload, BINDING_KEY: binding_of(call)},
+            ttl_seconds=self.ttl_seconds,
+        )
+        if not created:
+            raise RuntimeError("the store refused a fresh request-state id")
+        return state_id
+
+    async def read(self, call: Call, state_id: str) -> Any:
+        """Read the payload; raise KeyError for missing state or a different call binding."""
+        record = await self.store.get(self.key(state_id))
+        if record is None or record.data.get(BINDING_KEY) != binding_of(call):
+            raise KeyError(state_id)
+        return record.data.get(PAYLOAD_KEY)
+
+    async def drop(self, state_id: str) -> None:
+        await self.store.delete(self.key(state_id))
+
+
+async def restore_request_state(call: Call, states: RequestStates) -> Failure | None:
+    """Replace the id the client sent with the state the handler left."""
+    if call.state is None:
+        return None
+    if not isinstance(call.state, str):
+        return Failure(FailureKind.INVALID_PARAMS, "invalid requestState")
+    try:
+        call.state = await states.read(call, call.state)
+    except KeyError:
+        return Failure(FailureKind.INVALID_PARAMS, "invalid requestState")
+    return None
+
+
+async def protect_request_state(
+    call: Call, outcome: NeedsInput, states: RequestStates
+) -> NeedsInput | Failure:
+    """Keep the handler's state and hand the client its id instead."""
+    if outcome.state is None:
+        return outcome
+    try:
+        state_id = await states.open(call, outcome.state)
+    except (RuntimeError, TypeError) as e:
+        return Failure(FailureKind.INTERNAL, f"requestState could not be stored: {e}")
+    return NeedsInput(requests=outcome.requests, state=state_id)
+
+
+__all__ = [
+    "DEFAULT_TTL_SECONDS",
+    "RequestStates",
+    "binding_of",
+    "protect_request_state",
+    "restore_request_state",
+]

aiohttp_tiny_mcp/schema.py

@@ -0,0 +1,131 @@
+"""JSON Schema generation and simplification for legacy clients."""
+
+from __future__ import annotations
+
+import copy
+import re
+from typing import Any
+
+from pydantic import BaseModel
+
+HEADER_ANNOTATION = "x-mcp-header"
+HEADER_NAME = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$")
+HeaderParam = tuple[tuple[str, ...], str]
+
+
+def json_schema(model: type[BaseModel]) -> dict:
+    """Generate JSON Schema 2020-12 with refs intact; legacy adapters simplify it separately."""
+    schema = model.model_json_schema()
+    schema.setdefault("$schema", "https://json-schema.org/draft/2020-12/schema")
+    schema.setdefault("type", "object")
+    schema.setdefault("properties", {})
+    return schema
+
+
+def inline_refs(schema: dict[str, Any]) -> dict[str, Any]:
+    """Inline refs for legacy clients. Recursive schemas are not supported."""
+    defs = schema.get("$defs", {})
+
+    def walk(node: Any) -> Any:
+        if isinstance(node, dict):
+            ref = node.get("$ref")
+            if isinstance(ref, str) and ref.startswith("#/$defs/"):
+                target = walk(copy.deepcopy(defs[ref.rsplit("/", 1)[1]]))
+                return {**target, **{k: v for k, v in node.items() if k != "$ref"}}
+            return {k: walk(v) for k, v in node.items() if k != "$defs"}
+        if isinstance(node, list):
+            return [walk(x) for x in node]
+        return node
+
+    return walk(schema)
+
+
+def header_params(schema: dict[str, Any]) -> tuple[HeaderParam, ...]:
+    """Return (property_path, header_suffix) pairs for reachable x-mcp-header declarations. Reject
+    invalid declarations at registration.
+    """
+    found: list[HeaderParam] = []
+    seen: set[str] = set()
+    definitions = schema.get("$defs", {})
+
+    def walk(node: Any, path: tuple[str, ...], active_refs: frozenset[str] = frozenset()) -> None:
+        if not isinstance(node, dict):
+            return
+        ref = node.get("$ref")
+        if isinstance(ref, str) and ref.startswith("#/$defs/") and ref not in active_refs:
+            target = definitions.get(ref.rsplit("/", 1)[1])
+            if isinstance(target, dict):
+                walk(
+                    {**target, **{key: value for key, value in node.items() if key != "$ref"}},
+                    path,
+                    active_refs | {ref},
+                )
+                return
+        properties = node.get("properties")
+        if not isinstance(properties, dict):
+            return
+        for property_name, property_schema in properties.items():
+            if not isinstance(property_schema, dict):
+                continue
+            property_path = (*path, property_name)
+            annotation = property_schema.get(HEADER_ANNOTATION)
+            if annotation is not None:
+                if (
+                    not isinstance(annotation, str)
+                    or not annotation
+                    or not HEADER_NAME.fullmatch(annotation)
+                ):
+                    raise ValueError(f"invalid x-mcp-header value at {'.'.join(property_path)}")
+                if property_schema.get("type") not in {"string", "integer", "boolean"}:
+                    raise ValueError(
+                        f"x-mcp-header requires string, integer, or boolean at "
+                        f"{'.'.join(property_path)}"
+                    )
+                folded = annotation.casefold()
+                if folded in seen:
+                    raise ValueError(f"duplicate x-mcp-header value: {annotation}")
+                seen.add(folded)
+                found.append((property_path, annotation))
+            walk(property_schema, property_path, active_refs)
+
+    walk(schema, ())
+    return tuple(found)
+
+
+def simplify_legacy_schema(schema: dict) -> dict | None:
+    """Best-effort schema simplification for pre-2026-07-28 clients."""
+    node = _simplify_node(inline_refs(schema))
+    if node is None or node.get("type") != "object":
+        return None
+    return node
+
+
+def _simplify_node(node: Any) -> Any:
+    if not isinstance(node, dict):
+        return node
+    node = {k: v for k, v in node.items() if k not in ("if", "then", "else")}
+    any_of = node.get("anyOf")
+    if isinstance(any_of, list):
+        branches = [
+            branch
+            for branch in any_of
+            if not (isinstance(branch, dict) and branch.get("type") == "null")
+        ]
+        if len(branches) != 1:
+            return None
+        branch = _simplify_node(branches[0])
+        if branch is None:
+            return None
+        node = {**{k: v for k, v in node.items() if k != "anyOf"}, **branch}
+    if "oneOf" in node:
+        return None
+    properties = node.get("properties")
+    if isinstance(properties, dict):
+        simplified = {}
+        for key, prop in properties.items():
+            prop = _simplify_node(prop)
+            if prop is None:
+                return None
+            simplified[key] = prop
+        node = {**node, "properties": simplified}
+    return node