From b3623a8060cf8168b18df86a6e457de808afebd4 Mon Sep 17 00:00:00 2001
From: Eric Niebler <eniebler@nvidia.com>
Date: Fri, 30 Oct 2026 21:23:19 +0000
Subject: [PATCH] nvexec sync_wait uses correct env when querying for
 completion scheduler

---
 include/nvexec/stream/common.cuh        |  2 +-
 include/nvexec/stream/schedule_from.cuh |  2 +-
 include/nvexec/stream/sync_wait.cuh     | 23 +++++++++--------------
 3 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/include/nvexec/stream/common.cuh b/include/nvexec/stream/common.cuh
index 1fc2fa9cd..0198c5c77 100644
--- a/include/nvexec/stream/common.cuh
+++ b/include/nvexec/stream/common.cuh
@@ -104,7 +104,7 @@ namespace nv::execution
 
     template <class Sender, class Env>
     concept stream_completing_sender =
-      sender<Sender>
+      sender_in<Sender, Env>
       && gpu_stream_scheduler<
         __result_of<get_completion_scheduler<set_value_t>, env_of_t<Sender>, Env>,
         Env>;
diff --git a/include/nvexec/stream/schedule_from.cuh b/include/nvexec/stream/schedule_from.cuh
index d3f20f72b..5f39f5072 100644
--- a/include/nvexec/stream/schedule_from.cuh
+++ b/include/nvexec/stream/schedule_from.cuh
@@ -184,7 +184,7 @@ namespace nv::execution
     struct transform_sender_for<STDEXEC::schedule_from_t>
     {
       template <class Env, class Sender>
-      auto operator()(Env const & env, __ignore, __ignore, Sender&& sndr) const
+      constexpr auto operator()(Env const & env, __ignore, __ignore, Sender&& sndr) const
       {
         if constexpr (stream_completing_sender<Sender, Env>)
         {
diff --git a/include/nvexec/stream/sync_wait.cuh b/include/nvexec/stream/sync_wait.cuh
index 5549e28b7..14502c777 100644
--- a/include/nvexec/stream/sync_wait.cuh
+++ b/include/nvexec/stream/sync_wait.cuh
@@ -179,24 +179,26 @@ namespace nv::execution::_strm
 
     struct sync_wait_t
     {
-      template <sender_in<env> Sender>
+      template <stream_completing_sender<env> Sender>
         requires(__count_of<set_value_t, Sender, env>::value == 1)
-      auto operator()(context ctx, Sender&& sndr) const  //
-        -> std::optional<sync_wait_result_t<Sender>>
+      auto operator()(Sender&& sndr) const -> std::optional<sync_wait_result_t<Sender>>
       {
         using state_t = _sync_wait::state<Sender>;
         state_t  state{};
         run_loop loop;
 
+        auto sched = STDEXEC::get_completion_scheduler<set_value_t>(STDEXEC::get_env(sndr),
+                                                                    env{loop.get_scheduler()});
+
         cudaError_t status  = cudaSuccess;
         auto        opstate = host_allocate<exit_opstate_t<Sender, receiver<Sender>>>(
           status,
-          ctx.pinned_resource_,
+          sched.ctx_.pinned_resource_,
           __emplace_from{[&]
                          {
                            return exit_opstate(static_cast<Sender&&>(sndr),
                                                receiver<Sender>{{}, &state, &loop},
-                                               ctx);
+                                               sched.ctx_);
                          }});
         if (status != cudaSuccess)
         {
@@ -222,15 +224,8 @@ namespace nv::execution::_strm
   }  // namespace _sync_wait
 
   template <>
-  struct apply_sender_for<sync_wait_t>
-  {
-    template <stream_completing_sender<STDEXEC::env<>> Sender>
-    auto operator()(Sender&& sndr) const
-    {
-      auto sched = get_completion_scheduler<set_value_t>(get_env(sndr), STDEXEC::env{});
-      return _sync_wait::sync_wait_t{}(sched.ctx_, static_cast<Sender&&>(sndr));
-    }
-  };
+  struct apply_sender_for<sync_wait_t> : _sync_wait::sync_wait_t
+  {};
 }  // namespace nv::execution::_strm
 
 namespace nvexec = nv::execution;