"""Tests for `BaseContext`.

`BaseContext` is composition over a `DispatchContext` - it forwards

`transport`/`cancel_requested`/`send_raw_request`/`notify`/`progress`

and adds `meta`. It must satisfy `Outbound` so `ClientPeer` can wrap it.

"""

from collections.abc import Mapping

from typing import Any

import anyio

import pytest

from mcp.shared.context import BaseContext

from mcp.shared.dispatcher import DispatchContext

from mcp.shared.peer import ClientPeer

from mcp.shared.transport_context import TransportContext

from .conftest import direct_pair, jsonrpc_pair

from .test_dispatcher import Recorder, echo_handlers, running_pair

DCtx = DispatchContext[TransportContext]

@pytest.mark.anyio

async def test_base_context_forwards_transport_and_cancel_requested():

captured: list[BaseContext[TransportContext]] = []

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

captured.append(bctx)

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

bctx = captured[0]

assert bctx.transport.kind == "direct"

assert isinstance(bctx.cancel_requested, anyio.Event)

assert bctx.can_send_request is True

assert bctx.meta is None

@pytest.mark.anyio

async def test_base_context_can_send_request_reflects_dispatch_context_closed_state():

"""`can_send_request` must track the dctx, not the static transport flag,

so it agrees with whether `send_raw_request` would raise."""

captured: list[BaseContext[TransportContext]] = []

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

captured.append(BaseContext(ctx))

return {}

async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

bctx = captured[0]

assert bctx.transport.can_send_request is True

assert bctx.can_send_request is False

@pytest.mark.anyio

async def test_base_context_send_raw_request_and_notify_forward_to_dispatch_context():

crec = Recorder()

c_req, c_notify = echo_handlers(crec)

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

sample = await bctx.send_raw_request("sampling/createMessage", {"x": 1})

await bctx.notify("notifications/message", {"level": "info"})

return {"sample": sample}

async with running_pair(

direct_pair,

server_on_request=server_on_request,

client_on_request=c_req,

client_on_notify=c_notify,

) as (client, *_):

with anyio.fail_after(5):

result = await client.send_raw_request("tools/call", None)

await crec.notified.wait()

assert crec.requests == [("sampling/createMessage", {"x": 1})]

assert crec.notifications == [("notifications/message", {"level": "info"})]

assert result["sample"] == {"echoed": "sampling/createMessage", "params": {"x": 1}}

@pytest.mark.anyio

async def test_base_context_report_progress_invokes_caller_on_progress():

received: list[tuple[float, float | None, str | None]] = []

async def on_progress(progress: float, total: float | None, message: str | None) -> None:

received.append((progress, total, message))

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

await bctx.report_progress(0.5, total=1.0, message="halfway")

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None, {"on_progress": on_progress})

assert received == [(0.5, 1.0, "halfway")]

@pytest.mark.anyio

async def test_base_context_satisfies_outbound_so_peer_mixin_works():

"""Wrapping a BaseContext in ClientPeer proves it satisfies Outbound structurally."""

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

await ClientPeer(bctx).ping()

return {}

crec = Recorder()

c_req, c_notify = echo_handlers(crec)

async with running_pair(

direct_pair, server_on_request=server_on_request, client_on_request=c_req, client_on_notify=c_notify

) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

assert crec.requests == [("ping", None)]

@pytest.mark.anyio

async def test_base_context_meta_holds_supplied_request_params_meta():

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx, meta={"progressToken": "abc"})

assert bctx.meta is not None and bctx.meta.get("progressToken") == "abc"

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)