trace simple websocket messages - python

Hello XRay Team,

im having a project in python which uses websockets library. now we want to use xray to get some insights and i have to instrument it. as a starting point, i decided to set up a sample project to get to know xray, but it seems i cannot make it work.

when a client connects, im getting an error:

> connection handler failed
> Traceback (most recent call last):
> File "/home/ec2-user/environment/projects/xray/websocket.py", line 25, in echo
> subsegment = xray_recorder.begin_subsegment('process_message')
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 313, in begin_subsegment
> return self._begin_subsegment_helper(name, namespace)
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 287, in _begin_subsegment_helper
> segment = self.current_segment()
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 271, in current_segment
> entity = self.get_trace_entity()
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 411, in get_trace_entity
> return self.context.get_trace_entity()
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/context.py", line 101, in get_trace_entity
> return self.handle_context_missing()
> File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/context.py", line 125, in handle_context_missing
> raise SegmentNotFoundException(MISSING_SEGMENT_MSG)
> aws_xray_sdk.core.exceptions.exceptions.SegmentNotFoundException: cannot find the current segment/subsegment, please make sure you have a segment open

```
#!/usr/bin/env python

import asyncio
from websockets.server import serve

from aws_xray_sdk.core.async_context import AsyncContext
from aws_xray_sdk.core import xray_recorder

# Configure the X-Ray recorder
xray_recorder.configure(
service='websocket',
context=AsyncContext(),
context_missing='RUNTIME_ERROR',
daemon_address='localhost:2000'
)

async def echo(websocket):
segment = xray_recorder.begin_segment('echo')

try:
async for message in websocket:
subsegment = xray_recorder.begin_subsegment('process_message')

try:
await websocket.send(message)
except Exception as e:
subsegment.add_exception(e)
finally:
xray_recorder.end_subsegment()

except Exception as e:
segment.add_exception(e)
finally:
xray_recorder.end_segment()

async def main():
async with serve(echo, "localhost", 8765):
await asyncio.Future() # run forever

asyncio.run(main())
```

what am i doing wrong?