refactor: split RequestContext between server and client (#1987)

This commit is contained in:
Marcelo Trylesinski 2026-02-03 14:35:07 +01:00 • committed by GitHub
commit b1f7eec3cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 210 additions and 212 deletions

View file

@ -88,13 +88,13 @@ async def test_client_capabilities_with_tasks():
# Define custom handlers to trigger capability building (never actually called)
async def my_list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
raise NotImplementedError
async def my_cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
raise NotImplementedError
@ -168,13 +168,13 @@ async def test_client_capabilities_auto_built_from_handlers():
# Define custom handlers (not defaults)
async def my_list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
raise NotImplementedError
async def my_cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
raise NotImplementedError
@ -249,7 +249,7 @@ async def test_client_capabilities_with_task_augmented_handlers():
# Define task-augmented handler
async def my_augmented_sampling_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:

View file

@ -113,7 +113,7 @@ async def test_client_handles_get_task_request(client_streams: ClientTestStreams
received_task_id: str | None = None
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
nonlocal received_task_id
@ -176,7 +176,7 @@ async def test_client_handles_get_task_result_request(client_streams: ClientTest
store = InMemoryTaskStore()
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)
@ -239,7 +239,7 @@ async def test_client_handles_list_tasks_request(client_streams: ClientTestStrea
store = InMemoryTaskStore()
async def list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> ListTasksResult | ErrorData:
cursor = params.cursor if params else None
@ -294,7 +294,7 @@ async def test_client_handles_cancel_task_request(client_streams: ClientTestStre
store = InMemoryTaskStore()
async def cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CancelTaskRequestParams,
) -> CancelTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -361,7 +361,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
background_tg: list[TaskGroup | None] = [None]
async def task_augmented_sampling_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -384,7 +384,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
return CreateTaskResult(task=task)
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -400,7 +400,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
)
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)
@ -505,7 +505,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
background_tg: list[TaskGroup | None] = [None]
async def task_augmented_elicitation_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult | ErrorData:
@ -524,7 +524,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
return CreateTaskResult(task=task)
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -540,7 +540,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
)
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)

View file

@ -53,7 +53,7 @@ def create_client_task_handlers(
task_complete_events: dict[str, Event] = {}
async def handle_augmented_elicitation(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -72,7 +72,7 @@ def create_client_task_handlers(
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
@ -89,7 +89,7 @@ def create_client_task_handlers(
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
@ -121,7 +121,7 @@ def create_sampling_task_handlers(
task_complete_events: dict[str, Event] = {}
async def handle_augmented_sampling(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -140,7 +140,7 @@ def create_sampling_task_handlers(
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
@ -157,7 +157,7 @@ def create_sampling_task_handlers(
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
@ -211,7 +211,7 @@ async def test_scenario1_normal_tool_normal_elicitation() -> None:
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()
@ -379,7 +379,7 @@ async def test_scenario3_task_augmented_tool_normal_elicitation() -> None:
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()