wait_for_message

Instructions

Implementation Reference

• mcp_chat/server.py:326-428 (handler)

  The primary handler function for the 'wait_for_message' MCP tool. It uses long-polling with asyncio.Queue to wait for incoming messages in a chat room, handles timeouts, validation, and cleanup. Registered via @mcp.tool() decorator, with schema inferred from type hints and docstring.

  async def wait_for_message(
      room_id: str, client_id: str, timeout: int = 60
  ) -> Dict[str, Any]:
      """Wait for a message in the chat room (long-polling).
  
      This tool blocks until a message is received or the timeout is reached.
      Use this after sending a message to wait for a response, or call it first
      to wait for an incoming message.
  
      Conversation flow:
      - If you sent the last message: wait_for_message to get response
      - If you're waiting for first contact: wait_for_message before sending
      - After receiving a message: send_message to respond, then wait_for_message again
  
      Args:
          room_id: The ID of the chat room to listen in
          client_id: Your client identifier (from enter_queue or join_room)
          timeout: Timeout in seconds (default: 60, max: 300)
  
      Returns:
          On message: {"message": "text", "sender": "name", "timestamp": "...", "message_id": "..."}
          On timeout: {"timeout": true, "message": "No message received"}
          On error: {"error": "error message"}
      """
      # Use the provided client_id
      connection_id = client_id
  
      # Validate timeout
      timeout = min(timeout, 300)  # Max 5 minutes
      timeout = max(timeout, 1)  # Min 1 second
  
      # Get user
      user = connections.get(connection_id)
      if not user:
          logger.error(f"User not found for client_id: {client_id}")
          logger.debug(f"Active connections: {list(connections.keys())}")
          return {"error": f"User not found. Invalid client_id: {client_id}"}
  
      # Get and validate room
      room = await room_manager.get_room(room_id)
      if not room:
          return {"error": "Room not found"}
  
      if not room.active:
          return {"error": "Chat has ended"}
  
      # Verify user is in the room
      if not room.has_user(user.user_id):
          return {"error": "You are not in this room"}
  
      # Create a message queue for this user if not exists
      if room_id not in message_queues:
          message_queues[room_id] = {}
  
      # Create queue with reasonable size limit
      message_queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue(maxsize=100)
      message_queues[room_id][user.user_id] = message_queue
  
      logger.info(
          f"User {user.name} waiting for messages in room {room_id} (timeout: {timeout}s)"
      )
  
      try:
          # Wait for a message with timeout
          message_data = await asyncio.wait_for(
              message_queue.get(), timeout=float(timeout)
          )
  
          logger.info(
              f"Message received for {user.name}: {message_data.get('content', '')[:50]}..."
          )
  
          return {
              "message": message_data["content"],
              "sender": message_data["sender_name"],
              "timestamp": message_data["timestamp"],
              "message_id": message_data["message_id"],
          }
  
      except asyncio.TimeoutError:
          logger.info(f"Timeout waiting for message for {user.name}")
          return {"timeout": True, "message": "No message received within timeout period"}
  
      except asyncio.CancelledError:
          # Client cancelled the request - this is normal behavior
          logger.info(f"Wait cancelled for {user.name}")
          # Re-raise to let the framework handle it properly
          raise
  
      except Exception as e:
          logger.error(f"Error in wait_for_message: {e}")
          return {"error": f"Unexpected error: {str(e)}"}
  
      finally:
          # Clean up queue registration
          if room_id in message_queues and user.user_id in message_queues[room_id]:
              del message_queues[room_id][user.user_id]
              # Clean up empty room entries
              if not message_queues[room_id]:
                  del message_queues[room_id]
              logger.info(f"Cleaned up message queue for {user.name}")