Fix context shutdown

Updates #342 to merge cleanly into master.

Fixes #340, #341