diff --git a/Project.toml b/Project.toml
index b69146f17..141b0d3b8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
-version = "1.16.2"
+version = "1.16.3"
 authors = ["Francis Gagnon"]
 
 [deps]
diff --git a/src/controller/transcription.jl b/src/controller/transcription.jl
index d18d4365d..be0148ccd 100644
--- a/src/controller/transcription.jl
+++ b/src/controller/transcription.jl
@@ -1264,7 +1264,6 @@ function con_nonlinprogeq!(
     nΔU, nX̂ = nu*Hc, nx̂*Hp
     f_threads = transcription.f_threads
     Ts, p = model.Ts, model.p
-    As, Cs_u = mpc.estim.As, mpc.estim.Cs_u
     nk = get_nk(model, transcription)
     D̂0 = mpc.D̂0
     X̂0_Z̃ = @views Z̃[(nΔU+1):(nΔU+nX̂)]
@@ -1281,19 +1280,18 @@ function con_nonlinprogeq!(
         x̂0next   = @views   X̂0[(1 + nx̂*(j-1)):(nx̂*j)]
         x̂0next_Z̃ = @views X̂0_Z̃[(1 + nx̂*(j-1)):(nx̂*j)]  
         ŝnext    = @views  geq[(1 + nx̂*(j-1)):(nx̂*j)]  
-        x0, xs              = @views x̂0[1:nx], x̂0[nx+1:end]
+        x0                  = @views x̂0[1:nx]
+        xsnext              = @views x̂0next[nx+1:end]
         x0next_Z̃, xsnext_Z̃  = @views x̂0next_Z̃[1:nx], x̂0next_Z̃[nx+1:end]
         sdnext, ssnext      = @views ŝnext[1:nx], ŝnext[nx+1:end]
         k1, k2              = @views k0[1:nx], k0[nx+1:2*nx]
         # ----------------- stochastic defects -----------------------------------------
-        xsnext = @views x̂0next[nx+1:end]
-        mul!(xsnext, As, xs)
+        fs!(x̂0next, mpc.estim, model, x̂0)
         ssnext .= @. xsnext - xsnext_Z̃
         # ----------------- deterministic defects --------------------------------------
         u0 = @views U0[(1 + nu*(j-1)):(nu*j)]
         û0 = @views Û0[(1 + nu*(j-1)):(nu*j)]
-        mul!(û0, Cs_u, xs)                 # ys_u(k) = Cs_u*xs(k)
-        û0 .+= u0                          #   û0(k) = u0(k) + ys_u(k)
+        f̂_input!(û0, mpc.estim, model, x̂0, u0)
         if f_threads || h < 1 || j < 2
             # we need to recompute k1 with multi-threading, even with h==1, since the 
             # last iteration (j-1) may not be executed (iterations are re-orderable)
@@ -1307,8 +1305,7 @@ function con_nonlinprogeq!(
         else
             u0next = @views U0[(1 + nu*j):(nu*(j+1))]
             û0next = @views Û0[(1 + nu*j):(nu*(j+1))]
-            mul!(û0next, Cs_u, xsnext_Z̃)      # ys_u(k+1) = Cs_u*xs(k+1)
-            û0next .+= u0next                 #   û0(k+1) = u0(k+1) + ys_u(k+1)
+            f̂_input!(û0next, mpc.estim, model, x̂0next_Z̃, u0next)
         end
         model.f!(k2, x0next_Z̃, û0next, d̂0next, p)
         sdnext .= @. x0 - x0next_Z̃ + 0.5*Ts*(k1 + k2)
diff --git a/src/estimator/execute.jl b/src/estimator/execute.jl
index 45cc50558..e1fa8b076 100644
--- a/src/estimator/execute.jl
+++ b/src/estimator/execute.jl
@@ -108,6 +108,31 @@ function f̂!(x̂0next, û0, k0, model::SimModel, As, Cs_u, f̂op, x̂op, x̂0,
     return nothing
 end
 
+#TODO: delete the following two generic functions and replace with linear eq. constraints
+
+"""
+    fs!(x̂0next, estim::StateEstimator, model::SimModel, x̂0) -> nothing
+
+State update function of the stochastic model only.
+"""
+function fs!(x̂0next, estim::StateEstimator, model::SimModel, x̂0)
+    xs, xsnext = @views x̂0[model.nx+1:end], x̂0next[model.nx+1:end]
+    mul!(xsnext, estim.As, xs)
+    return nothing
+end
+
+@doc raw"""
+    f̂_input!(û0, estim::StateEstimator, model::SimModel, x̂0, u0) -> nothing
+
+Compute the disturbed input of the augmented model ``\mathbf{û_0}`` from `x̂0` and `u0`.
+"""
+function f̂_input!(û0, estim::StateEstimator, model::SimModel, x̂0, u0)
+    xs = @views x̂0[model.nx+1:end]
+    mul!(û0, estim.Cs_u, xs)      # ys_u = Cs_u*xs
+    û0 .+= u0                     # û0 = u0 + ys_u  
+    return nothing
+end
+
 @doc raw"""
     ĥ!(ŷ0, estim::StateEstimator, model::SimModel, x̂0, d0) -> nothing
 
@@ -141,7 +166,6 @@ function ĥ!(ŷ0, model::SimModel, Cs_y, x̂0, d0)
     return nothing
 end
 
-
 @doc raw"""
     initstate!(estim::StateEstimator, u, ym, d=[]) -> x̂
 
diff --git a/src/estimator/internal_model.jl b/src/estimator/internal_model.jl
index a73466cf5..0a5e04524 100644
--- a/src/estimator/internal_model.jl
+++ b/src/estimator/internal_model.jl
@@ -180,6 +180,22 @@ function f̂!(x̂0next, _ , k0, estim::InternalModel, model::NonLinModel, x̂0,
     return nothing
 end
 
+#TODO: delete the following 2 generic functions and replace with linear eq. constraints. 
+
+"""
+    fs!(x̂0next, estim::InternalModel, model::SimModel, _ ) -> nothing
+
+Does nothing since [`InternalModel`](@ref) does not augment the state vector.
+"""
+fs!( _ , ::InternalModel, ::SimModel, _ ) = nothing
+
+@doc raw"""
+    f̂_input!(û0, estim::InternalModel, model::SimModel, x̂0, u0) -> nothing
+
+Compute `û0 .= u0` since [`InternalModel`](@ref) does not augment the state vector.
+"""
+f̂_input!(û0, ::InternalModel, ::SimModel, _ , u0) = (û0 .= u0; nothing)
+
 @doc raw"""
     ĥ!(ŷ0, estim::InternalModel, model::NonLinModel, x̂0, d0)
 
diff --git a/test/3_test_predictive_control.jl b/test/3_test_predictive_control.jl
index e8f413966..c5f09dbce 100644
--- a/test/3_test_predictive_control.jl
+++ b/test/3_test_predictive_control.jl
@@ -989,7 +989,7 @@ end
     @test u ≈ [1.0] atol=5e-2
 
     transcription = TrapezoidalCollocation(1)    
-    nmpc5_1 = NonLinMPC(nonlinmodel_c; Nwt=[0], Hp=100, Hc=1, transcription)
+    nmpc5_1 = NonLinMPC(InternalModel(nonlinmodel_c); Nwt=[0], Hp=100, Hc=1, transcription)
     preparestate!(nmpc5_1, [0.0])
     u = moveinput!(nmpc5_1, [1/0.001])
     @test u ≈ [1.0] atol=5e-2