content: """
  Tests for TaskMessageQueue and InMemoryTaskMessageQueue.
  """
  
  from datetime import datetime, timezone
  
  import anyio
  import pytest
  
  from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
  from mcp.shared.experimental.tasks.resolver import Resolver
  from mcp.types import JSONRPCNotification, JSONRPCRequest
  
  
  @pytest.fixture
  def queue() -> InMemoryTaskMessageQueue:
      return InMemoryTaskMessageQueue()
  
  
  def make_request(id: int = 1, method: str = "test/method") -> JSONRPCRequest:
      return JSONRPCRequest(jsonrpc="2.0", id=id, method=method)
  
  
  def make_notification(method: str = "test/notify") -> JSONRPCNotification:
      return JSONRPCNotification(jsonrpc="2.0", method=method)
  
  
  class TestInMemoryTaskMessageQueue:
      @pytest.mark.anyio
      async def test_enqueue_and_dequeue(self, queue: InMemoryTaskMessageQueue) -> None:
          """Test basic enqueue and dequeue operations."""
          task_id = "task-1"
          msg = QueuedMessage(type="request", message=make_request())
  
          await queue.enqueue(task_id, msg)
          result = await queue.dequeue(task_id)
  
          assert result is not None
          assert result.type == "request"
          assert result.message.method == "test/method"
  
      @pytest.mark.anyio
      async def test_dequeue_empty_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
          """Dequeue from empty queue returns None."""
          result = await queue.dequeue("nonexistent-task")
          assert result is None
  
      @pytest.mark.anyio
      async def test_fifo_ordering(self, queue: InMemoryTaskMessageQueue) -> None:
          """Messages are dequeued in FIFO order."""
          task_id = "task-1"
  
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1, "first")))
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2, "second")))
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3, "third")))
  
          msg1 = await queue.dequeue(task_id)
          msg2 = await queue.dequeue(task_id)
          msg3 = await queue.dequeue(task_id)
  
          assert msg1 is not None and msg1.message.method == "first"
          assert msg2 is not None and msg2.message.method == "second"
          assert msg3 is not None and msg3.message.method == "third"
  
      @pytest.mark.anyio
      async def test_separate_queues_per_task(self, queue: InMemoryTaskMessageQueue) -> None:
          """Each task has its own queue."""
          await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1, "task1-msg")))
          await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2, "task2-msg")))
  
          msg1 = await queue.dequeue("task-1")
          msg2 = await queue.dequeue("task-2")
  
          assert msg1 is not None and msg1.message.method == "task1-msg"
          assert msg2 is not None and msg2.message.method == "task2-msg"
  
      @pytest.mark.anyio
      async def test_peek_does_not_remove(self, queue: InMemoryTaskMessageQueue) -> None:
          """Peek returns message without removing it."""
          task_id = "task-1"
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
  
          peeked = await queue.peek(task_id)
          dequeued = await queue.dequeue(task_id)
  
          assert peeked is not None
          assert dequeued is not None
          assert isinstance(peeked.message, JSONRPCRequest)
          assert isinstance(dequeued.message, JSONRPCRequest)
          assert peeked.message.id == dequeued.message.id
  
      @pytest.mark.anyio
      async def test_is_empty(self, queue: InMemoryTaskMessageQueue) -> None:
          """Test is_empty method."""
          task_id = "task-1"
  
          assert await queue.is_empty(task_id) is True
  
          await queue.enqueue(task_id, QueuedMessage(type="notification", message=make_notification()))
          assert await queue.is_empty(task_id) is False
  
          await queue.dequeue(task_id)
          assert await queue.is_empty(task_id) is True
  
      @pytest.mark.anyio
      async def test_clear_returns_all_messages(self, queue: InMemoryTaskMessageQueue) -> None:
          """Clear removes and returns all messages."""
          task_id = "task-1"
  
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1)))
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2)))
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3)))
  
          messages = await queue.clear(task_id)
  
          assert len(messages) == 3
          assert await queue.is_empty(task_id) is True
  
      @pytest.mark.anyio
      async def test_clear_empty_queue(self, queue: InMemoryTaskMessageQueue) -> None:
          """Clear on empty queue returns empty list."""
          messages = await queue.clear("nonexistent")
          assert messages == []
  
      @pytest.mark.anyio
      async def test_notification_messages(self, queue: InMemoryTaskMessageQueue) -> None:
          """Test queuing notification messages."""
          task_id = "task-1"
          msg = QueuedMessage(type="notification", message=make_notification("log/message"))
  
          await queue.enqueue(task_id, msg)
          result = await queue.dequeue(task_id)
  
          assert result is not None
          assert result.type == "notification"
          assert result.message.method == "log/message"
  
      @pytest.mark.anyio
      async def test_message_timestamp(self, queue: InMemoryTaskMessageQueue) -> None:
          """Messages have timestamps."""
          before = datetime.now(timezone.utc)
          msg = QueuedMessage(type="request", message=make_request())
          after = datetime.now(timezone.utc)
  
          assert before <= msg.timestamp <= after
  
      @pytest.mark.anyio
      async def test_message_with_resolver(self, queue: InMemoryTaskMessageQueue) -> None:
          """Messages can have resolvers."""
          task_id = "task-1"
          resolver: Resolver[dict[str, str]] = Resolver()
  
          msg = QueuedMessage(
              type="request",
              message=make_request(),
              resolver=resolver,
              original_request_id=42,
          )
  
          await queue.enqueue(task_id, msg)
          result = await queue.dequeue(task_id)
  
          assert result is not None
          assert result.resolver is resolver
          assert result.original_request_id == 42
  
      @pytest.mark.anyio
      async def test_cleanup_specific_task(self, queue: InMemoryTaskMessageQueue) -> None:
          """Cleanup removes specific task's data."""
          await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
          await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
  
          queue.cleanup("task-1")
  
          assert await queue.is_empty("task-1") is True
          assert await queue.is_empty("task-2") is False
  
      @pytest.mark.anyio
      async def test_cleanup_all(self, queue: InMemoryTaskMessageQueue) -> None:
          """Cleanup without task_id removes all data."""
          await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
          await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
  
          queue.cleanup()
  
          assert await queue.is_empty("task-1") is True
          assert await queue.is_empty("task-2") is True
  
      @pytest.mark.anyio
      async def test_wait_for_message_returns_immediately_if_message_exists(
          self, queue: InMemoryTaskMessageQueue
      ) -> None:
          """wait_for_message returns immediately if queue not empty."""
          task_id = "task-1"
          await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
  
          # Should return immediately, not block
          with anyio.fail_after(1):
              await queue.wait_for_message(task_id)
  
      @pytest.mark.anyio
      async def test_wait_for_message_blocks_until_message(self, queue: InMemoryTaskMessageQueue) -> None:
          """wait_for_message blocks until a message is enqueued."""
          task_id = "task-1"
          received = False
          waiter_started = anyio.Event()
  
          async def enqueue_when_ready() -> None:
              # Wait until the waiter has started before enqueueing
              await waiter_started.wait()
              await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
  
          async def wait_for_msg() -> None:
              nonlocal received
              # Signal that we're about to start waiting
              waiter_started.set()
              await queue.wait_for_message(task_id)
              received = True
  
          async with anyio.create_task_group() as tg:
              tg.start_soon(wait_for_msg)
              tg.start_soon(enqueue_when_ready)
  
          assert received is True
  
      @pytest.mark.anyio
      async def test_notify_message_available_wakes_waiter(self, queue: InMemoryTaskMessageQueue) -> None:
          """notify_message_available wakes up waiting coroutines."""
          task_id = "task-1"
          notified = False
          waiter_started = anyio.Event()
  
          async def notify_when_ready() -> None:
              # Wait until the waiter has started before notifying
              await waiter_started.wait()
              await queue.notify_message_available(task_id)
  
          async def wait_for_notification() -> None:
              nonlocal notified
              # Signal that we're about to start waiting
              waiter_started.set()
              await queue.wait_for_message(task_id)
              notified = True
  
          async with anyio.create_task_group() as tg:
              tg.start_soon(wait_for_notification)
              tg.start_soon(notify_when_ready)
  
          assert notified is True
  
      @pytest.mark.anyio
      async def test_peek_empty_queue_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
          """Peek on empty queue returns None."""
          result = await queue.peek("nonexistent-task")
          assert result is None
  
      @pytest.mark.anyio
      async def test_wait_for_message_double_check_race_condition(self, queue: InMemoryTaskMessageQueue) -> None:
          """wait_for_message returns early if message arrives after event creation but before wait."""
          task_id = "task-1"
  
          # To test the double-check path (lines 223-225), we need a message to arrive
          # after the event is created (line 220) but before event.wait() (line 228).
          # We simulate this by injecting a message before is_empty is called the second time.
  
          original_is_empty = queue.is_empty
          call_count = 0
  
          async def is_empty_with_injection(tid: str) -> bool:
              nonlocal call_count
              call_count += 1
              if call_count == 2 and tid == task_id:
                  # Before second check, inject a message - this simulates a message
                  # arriving between event creation and the double-check
                  queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())]
              return await original_is_empty(tid)
  
          queue.is_empty = is_empty_with_injection  # type: ignore[method-assign]
  
          # Should return immediately due to double-check finding the message
          with anyio.fail_after(1):
              await queue.wait_for_message(task_id)
  
  
  class TestResolver:
      @pytest.mark.anyio
      async def test_set_result_and_wait(self) -> None:
          """Test basic set_result and wait flow."""
          resolver: Resolver[str] = Resolver()
  
          resolver.set_result("hello")
          result = await resolver.wait()
  
          assert result == "hello"
          assert resolver.done()
  
      @pytest.mark.anyio
      async def test_set_exception_and_wait(self) -> None:
          """Test set_exception raises on wait."""
          resolver: Resolver[str] = Resolver()
  
          resolver.set_exception(ValueError("test error"))
  
          with pytest.raises(ValueError, match="test error"):
              await resolver.wait()
  
          assert resolver.done()
  
      @pytest.mark.anyio
      async def test_set_result_when_already_completed_raises(self) -> None:
          """Test that set_result raises if resolver already completed."""
          resolver: Resolver[str] = Resolver()
          resolver.set_result("first")
  
          with pytest.raises(RuntimeError, match="already completed"):
              resolver.set_result("second")
  
      @pytest.mark.anyio
      async def test_set_exception_when_already_completed_raises(self) -> None:
          """Test that set_exception raises if resolver already completed."""
          resolver: Resolver[str] = Resolver()
          resolver.set_result("done")
  
          with pytest.raises(RuntimeError, match="already completed"):
              resolver.set_exception(ValueError("too late"))
  
      @pytest.mark.anyio
      async def test_done_returns_false_before_completion(self) -> None:
          """Test done() returns False before any result is set."""
          resolver: Resolver[str] = Resolver()
          assert resolver.done() is False
  