diff --git a/internal/e2e/profiler/profiler.go b/internal/e2e/profiler/profiler.go
index a444ad7e..f7a10ea7 100644
--- a/internal/e2e/profiler/profiler.go
+++ b/internal/e2e/profiler/profiler.go
@@ -58,8 +58,9 @@ func main() {
 	p := wzprof.ProfilingFor(guestBin)
 	cpu := p.CPUProfiler()
 	mem := p.MemoryProfiler()
-	ctx := context.WithValue(context.Background(),
-		experimental.FunctionListenerFactoryKey{},
+
+	ctx := experimental.WithFunctionListenerFactory(
+		context.Background(),
 		experimental.MultiFunctionListenerFactory(
 			wzprof.Sample(sampleRate, cpu),
 			wzprof.Sample(sampleRate, mem),