ApiGateway Websocket API: async python Lambda function

Hi,

I'm having trouble getting ApiGateway Websocket API to work using an async python Lambda function.

When I try to connect over wscat, I get "error: Unexpected server response: 502". When I look at logs for the OnConnect function, I see "Runtime.MarshalError: Unable to marshal response: Object of type coroutine is not JSON serializable"

Is it not possible to invoke an async lambda for a websocket API? Or am I doing something wrong?

Any help much appreciated! TIA

Here is my simple lambda function ws_service.on_connect_function:

async def connect_function(event, context):  
    return {  
        "statusCode": 200,  
    }  

And here's the portion of my template that sets up the ws API and the $connect function integration:

  Ss2WebsocketApi:  
    Type: AWS::ApiGatewayV2::Api  
    Properties:  
      Name: Ss2WebsocketApi  
      ProtocolType: WEBSOCKET  
      RouteSelectionExpression: "$request.body.action"  
  ConnectRoute:  
    Type: AWS::ApiGatewayV2::Route  
    Properties:  
      ApiId: !Ref Ss2WebsocketApi  
      RouteKey: $connect  
      AuthorizationType: NONE  
      OperationName: ConnectRoute  
      Target: !Join  
        - "/"  
        - - "integrations"  
          - !Ref ConnectInteg  
  ConnectInteg:  
    Type: AWS::ApiGatewayV2::Integration  
    Properties:  
      ApiId: !Ref Ss2WebsocketApi  
      Description: Connect Integration  
      IntegrationType: AWS_PROXY  
      IntegrationUri:  
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnConnectFunction.Arn}/invocations  
  OnConnectFunction:  
    Type: AWS::Serverless::Function  
    Properties:  
      Handler: ws_service.connect_function  
  OnConnectPermission:  
    Type: AWS::Lambda::Permission  
    DependsOn:  
      - Ss2WebsocketApi  
      - OnConnectFunction  
    Properties:  
      Action: lambda:InvokeFunction  
      FunctionName: !Ref OnConnectFunction  
      Principal: apigateway.amazonaws.com