Query highlighting is broken in async context

[azmeuk]

Query highlighting is ok for regular contexts, but not in async contexts. This can be seen in the demo.

The fix should be simple, but I think we should make a test for this to prevent future regressions.