Use supports_callback_vjp trait and preserve FakeIntegrator types

Replaces the hardcoded `Union{ReverseDiffVJP, EnzymeVJP, ReactantVJP,
MooncakeVJP}` checks in `_setup_reverse_callbacks` and `setup_w_wp` with
a `supports_callback_vjp(autojacvec)` trait defined alongside
`supports_structured_vjp` in sensitivity_algorithms.jl. Future VJP
backends opt in by adding a single method instead of threading a new
type through two Union annotations.

Also fixes `get_FakeIntegrator(::MooncakeVJP, ...)` to pass `u` and `p`
through unchanged (matching the EnzymeVJP / ReactantVJP methods) instead
of using `[x for x in u]` / `[x for x in p]`. The comprehension is only
needed by the ReverseDiffVJP path to hand ReverseDiff a fresh tracked
container; Mooncake's pullback cache is fed by the dedicated
`cb_state_fn` / `cb_param_fn` closures in `get_cb_diffcaches`, so this
dispatch only fires at the plain `w(y, y, integrator.p, integrator.t)`
call site where preserving the original container type (e.g.
ComponentArray, MVector) matters.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

src/callback_tracking.jl

@@ -283,12 +283,12 @@ function _setup_reverse_callbacks(
     # if save_positions = [1,0] the gradient contribution is added before, and in principle we would need to correct the adjoint state again. Therefore,
 
     cb.save_positions == [1, 0] && error("save_positions=[1,0] is currently not supported.")
-    # Callbacks require ReverseDiffVJP, EnzymeVJP, ReactantVJP, or MooncakeVJP for
-    # their own VJP computations. The ODE adjoint may use a different autojacvec
-    # (even numerical/false), but the callback affect functions (CallbackAffectWrapper)
-    # are separate and typically work with ReverseDiff even when the ODE function doesn't.
-    cb_autojacvec = if sensealg.autojacvec isa
-            Union{ReverseDiffVJP, EnzymeVJP, ReactantVJP, MooncakeVJP}
+    # Callback affect functions (CallbackAffectWrapper) are traced separately
+    # from the ODE function, so the ODE adjoint may use a different autojacvec
+    # (even numerical/false) while the callback path uses its own compatible
+    # backend. The `supports_callback_vjp` trait marks which VJP backends have
+    # a dedicated callback path; anything else falls back to ReverseDiffVJP.
+    cb_autojacvec = if supports_callback_vjp(sensealg.autojacvec)
         sensealg.autojacvec
     else
         @warn "autojacvec=$(sensealg.autojacvec) is not compatible with callbacks, using ReverseDiffVJP() for callback VJPs"
@@ -507,9 +507,10 @@ end
 
 function setup_w_wp(
         cb::Union{DiscreteCallback, ContinuousCallback, VectorContinuousCallback},
-        autojacvec::Union{ReverseDiffVJP, EnzymeVJP, ReactantVJP, MooncakeVJP}, pos_neg,
-        event_idx, tprev
+        autojacvec, pos_neg, event_idx, tprev
     )
+    supports_callback_vjp(autojacvec) ||
+        error("setup_w_wp called with a VJP backend that does not support callbacks: $(autojacvec). This is an internal error — the callback path should have redirected to a compatible backend in `_setup_reverse_callbacks`.")
     w = CallbackAffectWrapper(cb, autojacvec, pos_neg, event_idx, tprev)
     wp = CallbackAffectPWrapper(cb, autojacvec, pos_neg, event_idx, tprev)
     return w, wp
@@ -520,9 +521,7 @@ function get_FakeIntegrator(autojacvec::ReverseDiffVJP, u, p, t, tprev)
 end
 get_FakeIntegrator(autojacvec::EnzymeVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
 get_FakeIntegrator(autojacvec::ReactantVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
-function get_FakeIntegrator(autojacvec::MooncakeVJP, u, p, t, tprev)
-    return FakeIntegrator([x for x in u], [x for x in p], t, tprev)
-end
+get_FakeIntegrator(autojacvec::MooncakeVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
 
 function _get_wp_paramjac_config(autojacvec::EnzymeVJP, _p, wp, y, __p, _t)
     return (zero(y), zero(_p), zero(_p), zero(_p), zero(y))

src/sensitivity_algorithms.jl

@@ -1696,6 +1696,23 @@ supports_structured_vjp(::ReverseDiffVJP) = false
 supports_structured_vjp(::Bool) = false
 supports_structured_vjp(::Nothing) = false
 
+"""
+    supports_callback_vjp(autojacvec) -> Bool
+
+Return `true` if the VJP backend can differentiate the affect functions of
+tracked callbacks (i.e. has a dedicated branch in `get_cb_diffcaches` and a
+`_vecjacobian!` path that works with `CallbackSensitivityFunction`).
+
+When `false`, the callback adjoint code falls back to `ReverseDiffVJP()` for
+its own VJP computations — the ODE adjoint may still use the user-requested
+backend, since the callback affect functions are traced separately.
+"""
+supports_callback_vjp(::ReverseDiffVJP) = true
+supports_callback_vjp(::EnzymeVJP) = true
+supports_callback_vjp(::ReactantVJP) = true
+supports_callback_vjp(::MooncakeVJP) = true
+supports_callback_vjp(::Any) = false
+
 """
 ```julia
 ForwardDiffOverAdjoint{A} <: AbstractSecondOrderSensitivityAlgorithm{nothing, true, nothing}