modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

perf: use deque for InMemoryTaskMessageQueue FIFO operations (#2165)

Browse source
This commit is contained in:
Giulio Leone 2026-03-05 16:44:33 +01:00 • committed by GitHub
commit b33c811675
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

9
src/mcp/shared/experimental/tasks/message_queue.py
View file

@ -12,6 +12,7 @@ This pattern enables:
""" """
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections import deque
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Any, Literal from typing import Any, Literal
@ -151,13 +152,13 @@ class InMemoryTaskMessageQueue(TaskMessageQueue):
""" """
def __init__(self) -> None: def __init__(self) -> None:
self._queues: dict[str, list[QueuedMessage]] = {} self._queues: dict[str, deque[QueuedMessage]] = {}
self._events: dict[str, anyio.Event] = {} self._events: dict[str, anyio.Event] = {}
def _get_queue(self, task_id: str) -> list[QueuedMessage]: def _get_queue(self, task_id: str) -> deque[QueuedMessage]:
"""Get or create the queue for a task.""" """Get or create the queue for a task."""
if task_id not in self._queues: if task_id not in self._queues:
self._queues[task_id] = [] self._queues[task_id] = deque()
return self._queues[task_id] return self._queues[task_id]
async def enqueue(self, task_id: str, message: QueuedMessage) -> None: async def enqueue(self, task_id: str, message: QueuedMessage) -> None:
@ -172,7 +173,7 @@ class InMemoryTaskMessageQueue(TaskMessageQueue):
queue = self._get_queue(task_id) queue = self._get_queue(task_id)
if not queue: if not queue:
return None return None
return queue.pop(0) return queue.popleft()
async def peek(self, task_id: str) -> QueuedMessage | None: async def peek(self, task_id: str) -> QueuedMessage | None:
"""Return the next message without removing it.""" """Return the next message without removing it."""

3
tests/experimental/tasks/test_message_queue.py
View file

@ -1,5 +1,6 @@
"""Tests for TaskMessageQueue and InMemoryTaskMessageQueue.""" """Tests for TaskMessageQueue and InMemoryTaskMessageQueue."""
from collections import deque
from datetime import datetime, timezone from datetime import datetime, timezone
import anyio import anyio
@ -270,7 +271,7 @@ class TestInMemoryTaskMessageQueue:
if call_count == 2 and tid == task_id: if call_count == 2 and tid == task_id:
# Before second check, inject a message - this simulates a message # Before second check, inject a message - this simulates a message
# arriving between event creation and the double-check # arriving between event creation and the double-check
queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())] queue._queues[task_id] = deque([QueuedMessage(type="request", message=make_request())])
return await original_is_empty(tid) return await original_is_empty(tid)
queue.is_empty = is_empty_with_injection # type: ignore[method-assign] queue.is_empty = is_empty_with_injection # type: ignore[method-assign]