From b3f5a0c51b66995cd6e5b15d0da81121e12e75c8 Mon Sep 17 00:00:00 2001
From: Joshua Terranova <terranovajoshua1@gmail.com>
Date: Mon, 10 Aug 2026 18:37:19 -0700
Subject: [PATCH] ENH: add Flight post-step callback across all phases (#758)

---
 rocketpy/simulation/flight.py        | 22 +++++++++++++++
 tests/unit/simulation/test_flight.py | 41 ++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py
index 55ca3486f..ed4f78164 100644
--- a/rocketpy/simulation/flight.py
+++ b/rocketpy/simulation/flight.py
@@ -79,6 +79,14 @@ class Flight:
         Name of the flight.
     Flight._controllers : list
         List of controllers to be used during simulation.
+    Flight.post_step_callback : callable, optional
+        Optional callback invoked once after every successful ODE solver
+        step for the entire simulation, including parachute descent.
+        Receives the ``Flight`` instance (``callback(flight)``). Use
+        ``flight.t`` and ``flight.y_sol`` for the current time and state.
+        Controllers stop being useful after parachute deployment; this
+        callback is the extension point for full-lifecycle observers
+        (e.g. ground-station / radio update simulation).
     Flight.max_time : int, float
         Maximum simulation time allowed. Refers to physical time
         being simulated, not time taken to run simulation.
@@ -506,6 +514,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-statements
         equations_of_motion="standard",
         ode_solver="LSODA",
         simulation_mode="6 DOF",
+        post_step_callback=None,
     ):
         """Run a trajectory simulation.
 
@@ -592,6 +601,12 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-statements
         simulation_mode : str, optional
             Simulation mode to use. Can be "6 DOF" for 6 degrees of freedom or
             "3 DOF" for 3 degrees of freedom. Default is "6 DOF".
+        post_step_callback : callable, optional
+            Callback invoked once after every successful ODE solver step for
+            the entire simulation, including parachute phases. Signature is
+            ``callback(flight)``, matching phase/node callbacks. Access the
+            current time and state via ``flight.t`` and ``flight.y_sol``.
+            Default is None.
         Returns
         -------
         None
@@ -600,6 +615,9 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-statements
         ----------
         .. [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html
         """
+        if post_step_callback is not None and not callable(post_step_callback):
+            raise TypeError("post_step_callback must be callable or None")
+
         # Save arguments
         self.env = environment
         self.rocket = rocket
@@ -626,6 +644,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-statements
         self.equations_of_motion = equations_of_motion
         self.simulation_mode = simulation_mode
         self.ode_solver = ode_solver
+        self.post_step_callback = post_step_callback
 
         # Controller initialization
         self.__init_controllers()
@@ -800,6 +819,9 @@ def __simulate(self, verbose):
                             self.sensors,
                             self.env,
                         )
+                    # Full-lifecycle observer (all phases, including parachute)
+                    if self.post_step_callback is not None:
+                        self.post_step_callback(self)
                     if self.__check_simulation_events(phase, phase_index, node_index):
                         break  # Stop if simulation termination event occurred
 
diff --git a/tests/unit/simulation/test_flight.py b/tests/unit/simulation/test_flight.py
index 9a3c54477..eca9d99fb 100644
--- a/tests/unit/simulation/test_flight.py
+++ b/tests/unit/simulation/test_flight.py
@@ -139,6 +139,47 @@ def test_get_controller_observed_variables(flight_calisto_air_brakes):
     assert len(obs_vars) == 0
 
 
+def test_post_step_callback_runs_before_and_after_apogee(
+    calisto_robust, example_plain_env
+):
+    """post_step_callback must fire across the full flight, including descent.
+
+    Controllers are not a substitute: air-brake fixtures often terminate at
+    apogee, and parachute phases are not meant to keep feeding actuators.
+    This callback is the full-lifecycle observer hook (issue #758).
+    """
+    callback_times = []
+
+    def record_step(flight):
+        callback_times.append(flight.t)
+
+    flight = Flight(
+        rocket=calisto_robust,
+        environment=example_plain_env,
+        rail_length=5.2,
+        inclination=85,
+        heading=0,
+        terminate_on_apogee=False,
+        post_step_callback=record_step,
+    )
+
+    assert callback_times, "post_step_callback was never invoked"
+    assert min(callback_times) < flight.apogee_time
+    assert max(callback_times) > flight.apogee_time
+    assert flight.t_final > flight.apogee_time
+
+
+def test_post_step_callback_must_be_callable(calisto, example_plain_env):
+    """Non-callable post_step_callback values are rejected at construction."""
+    with pytest.raises(TypeError, match="post_step_callback"):
+        Flight(
+            rocket=calisto,
+            environment=example_plain_env,
+            rail_length=5.2,
+            post_step_callback="not-callable",
+        )
+
+
 def test_initial_stability_margin(flight_calisto_custom_wind):
     """Test the initial_stability_margin method of the Flight class.