diff --git a/poll.go b/poll.go
index 195a261..e214952 100644
--- a/poll.go
+++ b/poll.go
@@ -279,7 +279,7 @@ func (h *Headscale) PollNetMapStream(
 	keepAliveChan chan []byte,
 	updateChan chan struct{},
 ) {
-	ctx := context.WithValue(context.Background(), machineNameContextKey, machine.Hostname)
+	ctx := context.WithValue(r.Context(), machineNameContextKey, machine.Hostname)
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()