From a79f4c2918b309cf4653f5cb110812cddaa214db Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 16 Dec 2024 19:25:42 +0530
Subject: [PATCH] fix: use `extra_args` in `generate_control_function`

---
 src/inputoutput.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/inputoutput.jl b/src/inputoutput.jl
index 1be1fbb8dd..6bdcac6dd4 100644
--- a/src/inputoutput.jl
+++ b/src/inputoutput.jl
@@ -252,10 +252,11 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     process = get_postprocess_fbody(sys)
     wrapped_arrays_vars = disturbance_argument ?
                           wrap_array_vars(
-        sys, rhss; dvs, ps, inputs, cachesyms = (disturbance_inputs,)) :
+        sys, rhss; dvs, ps, inputs, extra_args = (disturbance_inputs,)) :
                           wrap_array_vars(sys, rhss; dvs, ps, inputs)
     f = build_function(rhss, args...; postprocess_fbody = process,
-        expression = Val{true}, wrap_code = wrap_mtkparameters(sys, false, 3) .∘
+        expression = Val{true}, wrap_code = wrap_mtkparameters(
+            sys, false, 3, Int(disturbance_argument) + 1) .∘
                                             wrapped_arrays_vars .∘
                                             wrap_parameter_dependencies(sys, false),
         kwargs...)