refactor: split RequestContext between server and client (#1987)
This commit is contained in:
parent
050aeb6456
commit
b1f7eec3cd
3 changed files with 210 additions and 212 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue