From 9fd1fec7cb127be6f8f158f0a4656b3920444a91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 11:37:54 +0000
Subject: [PATCH 1/6] Preserve scenario inputs and refresh marginal-rate
 calculations

---
 ...fix-scenario-modifier-inputs-2666.fixed.md |   1 +
 docs/book/usage/scenarios.md                  |  24 +
 policyengine_uk/simulation.py                 |  79 ++--
 .../tests/test_scenario_modifier_inputs.py    | 418 ++++++++++++++++++
 .../variables/household/marginal_tax_rate.py  |  55 ++-
 .../marginal_tax_rate_wrt_employer_cost.py    |  57 +--
 6 files changed, 548 insertions(+), 86 deletions(-)
 create mode 100644 changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
 create mode 100644 policyengine_uk/tests/test_scenario_modifier_inputs.py

diff --git a/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
new file mode 100644
index 0000000000..12873d9ef7
--- /dev/null
+++ b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
@@ -0,0 +1 @@
+Keep inputs stored by construction-time scenario modifiers through resets and marginal-rate calculations, and clear their stale calculated values so labour supply responses reflect the completed scenario.
diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md
index 706772f6d5..388bfadc42 100644
--- a/docs/book/usage/scenarios.md
+++ b/docs/book/usage/scenarios.md
@@ -159,6 +159,30 @@ A sum keeps the timing of each part's modifier. A simulation runs the modifiers
 
 Beyond simple parameter changes, you sometimes need scenarios that make structural modifications to how the tax-benefit system works. Let's walk through how to create these, using examples similar to those in PolicyEngine UK's scenarios module.
 
+### Setting inputs in a simulation modifier
+
+When you pass a scenario to `Simulation` or `Microsimulation`, its modifier runs
+after data load by default. Values it stores with `set_input` become inputs in
+`simulation.input_variables`. They survive `reset_calculations`, marginal-rate
+branches and `apply_dynamics`, and dataset exports include them as inputs.
+The simulation drops calculated values after the modifier finishes, so later
+calculations use the inputs the modifier left behind. Calculating a variable
+inside the modifier does not make that variable an input.
+Employment marginal-rate calculations also take the completed scenario's
+inputs when the modifier calculated a marginal rate before changing an input.
+
+An input remains fixed when a marginal-rate calculation changes earnings. A
+modifier that adds a fixed amount to Child Benefit therefore keeps that amount
+through the earnings step. For a policy that should vary with earnings, use
+parameter changes or a formula reform; setting its calculated value once cannot
+represent that dependence. This cleanup happens when the constructor runs the
+modifier. Calling a modifier directly, or using `Scenario.apply` on an existing
+simulation, retains those operations' existing behavior.
+
+Modifiers with `applied_before_data_load=True` run before the simulation has its
+data. They can change parameters or model structure at that stage; modifiers
+that read or set household inputs belong after data load.
+
 ### Creating a two-child limit repeal scenario
 
 The two-child limit restricts Universal Credit child elements to the first two children. Here's how to build a scenario that removes this restriction:
diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
index f1a359bde9..efab5cd213 100644
--- a/policyengine_uk/simulation.py
+++ b/policyengine_uk/simulation.py
@@ -2,6 +2,7 @@
 import os
 import copy
 import warnings
+from contextlib import contextmanager
 from importlib.metadata import version
 from typing import Dict, Optional, Union, Type, List
 
@@ -57,6 +58,7 @@
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.data_source import built_from_data
 from policyengine_uk.utils.input_dataframe_export import input_dataframe
+from policyengine_uk.utils.measurement_branches import drop_calculated_values
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
     copy_model_values,
@@ -113,7 +115,7 @@ def _canonical_parameter_paths(changes: dict) -> dict:
 
 
 class _InputKeyLog(set):
-    """A simulation's record of its inputs while Simulation.apply_reform runs.
+    """A simulation's record of inputs stored by a reform or scenario modifier.
 
     policyengine-core records each input in the simulation's
     ``_user_input_keys`` as (variable, branch, period). ``Holder._set`` adds
@@ -124,10 +126,10 @@ class _InputKeyLog(set):
     variable's end, one for a neutralized variable, or one for a sub-period
     that a ``set_input`` helper skips because it already holds an input.
     This set also adds each key to every set in ``logs``: one for each
-    reform being applied to the simulation, the innermost last.
+    operation being applied to the simulation, the innermost last.
 
     A copy or pickle of it is a plain set, so a simulation deep-copied or
-    pickled while a reform runs logs nothing. A shallow copy of the
+    pickled while an operation runs logs nothing. A shallow copy of the
     simulation shares its record, as it shares its holders.
     """
 
@@ -144,11 +146,11 @@ def __reduce__(self):
         return set, (list(self),)
 
 
-# The logs (see _InputKeyLog) of the reforms Simulation.apply_reform is
-# applying, by the id of the simulation it applies them to, while it does.
+# The logs (see _InputKeyLog) of the reforms or scenario modifiers running,
+# by the id of the simulation they apply to, while they do.
 # Kept outside the simulation, because core's clone copies its __dict__:
 # a clone or branch made while a reform runs must not log for it.
-_reform_logs: Dict[int, list] = {}
+_input_key_logs: Dict[int, list] = {}
 
 
 # Cache for fully-loaded, uprated, enum-pre-encoded multi-year datasets keyed
@@ -433,7 +435,9 @@ def __init__(
             if self._construction_reform_state is not None:
                 self._construction_reform_state.refresh_framework(self)
             if after_data_load.simulation_modifier is not None:
-                after_data_load.simulation_modifier(self)
+                with self._preserve_stored_inputs():
+                    after_data_load.simulation_modifier(self)
+                drop_calculated_values(self)
         self._refresh_capital_gains_projection()
         self._capture_cgt_supplied_inputs()
         self._cgt_warning_enabled = True
@@ -456,48 +460,34 @@ def _refresh_capital_gains_projection(self) -> None:
         add_capital_gains_projection(self.tax_benefit_system.parameters, projection)
         self.tax_benefit_system.reset_parameter_caches()
 
-    def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
-        """Apply a runtime reform while keeping generated CGT assumptions coherent."""
-        # Reforms that set inputs (disable_simulated_benefits, adjust_budgets)
-        # reach the simulation through its tax-benefit system's link. A clone
-        # or branch may share its parent's system, so link the system to this
-        # simulation while the reform is applied, then restore the link.
-        system = self.tax_benefit_system
-        linked = getattr(system, "simulation", None)
-        system.simulation = self
-        # Log the inputs the reform stores (see _InputKeyLog and the
-        # _user_input_keys property). Within another reform applied to this
-        # simulation (a tuple's parts, or a reform that applies one), log
-        # them for that one too.
-        logs = _reform_logs.setdefault(id(self), [])
+    @contextmanager
+    def _preserve_stored_inputs(self):
+        """Name inputs an operation stores and retains, for later resets."""
+        logs = _input_key_logs.setdefault(id(self), [])
         stored = set()
         logs.append(stored)
         self._user_input_keys = getattr(self, "_user_input_keys", None) or set()
         try:
-            super().apply_reform(reform)
-            self._refresh_capital_gains_projection()
+            yield
         finally:
             logs.pop()
             if not logs:
-                del _reform_logs[id(self)]
-            # A plain set again, unless an enclosing reform is still logging.
+                del _input_key_logs[id(self)]
+            # A plain set again, unless an enclosing operation is still logging.
             self._user_input_keys = set(self._user_input_keys)
-            if linked is not None:
-                system.simulation = linked
         # input_variables is taken when the simulation is built.
         # reset_calculations deletes the values of every variable it does not
         # name (as core's derivative did before policyengine-core 3.33.0, from
         # which it keeps every input), and the marginal rate formulas those of
-        # every such variable that has a formula. A reform applied afterwards
-        # (a Reform class passed as reform=, or one applied to a clone or
-        # branch) can store inputs: new ones, or ones that replace an input
-        # set since the build. Name every variable with an input the reform
+        # every such variable that has a formula. A reform or scenario modifier
+        # applied afterwards can store inputs: new ones, or ones that replace
+        # an input set since the build. Name every variable with an input it
         # stored that this simulation still holds, under a branch name it
         # reads and a reset deletes, or a reset would delete what the reform
         # stored and recalculate the formula. input_variables names whole
         # variables, so a reset then keeps every input of the variable, any
-        # set since the build included. A reform that raises names nothing it
-        # stored itself (a reform it applied first has named its own). Assign
+        # set since the build included. An operation that raises names nothing
+        # it stored itself (a reform it applied first has named its own). Assign
         # a new list: a clone or branch shares its parent's.
         input_variables = getattr(self, "input_variables", None)
         if input_variables is not None:
@@ -511,13 +501,30 @@ def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
             if added:
                 self.input_variables = list(input_variables) + added
 
+    def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
+        """Apply a runtime reform while keeping generated CGT assumptions coherent."""
+        # Reforms that set inputs (disable_simulated_benefits, adjust_budgets)
+        # reach the simulation through its tax-benefit system's link. A clone
+        # or branch may share its parent's system, so link the system to this
+        # simulation while the reform is applied, then restore the link.
+        system = self.tax_benefit_system
+        linked = getattr(system, "simulation", None)
+        system.simulation = self
+        try:
+            with self._preserve_stored_inputs():
+                super().apply_reform(reform)
+                self._refresh_capital_gains_projection()
+        finally:
+            if linked is not None:
+                system.simulation = linked
+
     @property
     def _user_input_keys(self) -> set:
         # policyengine-core's record of the simulation's inputs, kept in
-        # __dict__ as core keeps it. While apply_reform applies a reform to
+        # __dict__ as core keeps it. While a reform or scenario modifier runs on
         # this simulation, a set assigned here becomes a log (see
         # _InputKeyLog), so a record core replaces mid-reform still logs: a
-        # branch's, when a reform applied to its parent rebuilds it, or one
+        # branch's, when an operation on its parent rebuilds it, or one
         # rebuilt when the system is rebound.
         try:
             return self.__dict__["_user_input_keys"]
@@ -526,7 +533,7 @@ def _user_input_keys(self) -> set:
 
     @_user_input_keys.setter
     def _user_input_keys(self, keys) -> None:
-        logs = _reform_logs.get(id(self))
+        logs = _input_key_logs.get(id(self))
         if logs and keys is not None:
             keys = _InputKeyLog(keys, logs)
         self.__dict__["_user_input_keys"] = keys
diff --git a/policyengine_uk/tests/test_scenario_modifier_inputs.py b/policyengine_uk/tests/test_scenario_modifier_inputs.py
new file mode 100644
index 0000000000..9fd686529a
--- /dev/null
+++ b/policyengine_uk/tests/test_scenario_modifier_inputs.py
@@ -0,0 +1,418 @@
+"""Construction modifiers retain explicit inputs and discard stale calculations.
+
+All situations are synthetic households. A fixed Child Benefit addition changes
+income but no marginal rate at these earnings, so it has only an income effect.
+The equivalent parameter change supplies an independent route to that outcome.
+"""
+
+import copy
+
+import numpy as np
+import pytest
+from policyengine_core.periods import period
+
+from policyengine_uk import Microsimulation, Simulation
+from policyengine_uk.model_api import Scenario
+
+pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
+
+YEAR = 2027
+CHILD_BENEFIT = "child_benefit"
+ELDEST_AMOUNT = "gov.hmrc.child_benefit.amount.eldest"
+RESPONSE = "employment_income_behavioral_response"
+INCOME_TERM = "income_elasticity_lsr"
+SUBSTITUTION_TERM = "substitution_elasticity_lsr"
+ELASTICITIES = {
+    "gov.simulation.labour_supply_responses.substitution_elasticity": {str(YEAR): 0.15},
+    "gov.simulation.labour_supply_responses.income_elasticity": {str(YEAR): -0.05},
+}
+
+
+def family(years=(YEAR,)):
+    return {
+        "people": {
+            "parent": {
+                "age": {year: 30 for year in years},
+                "gender": {year: "FEMALE" for year in years},
+                "employment_income": {year: 25_000 for year in years},
+                "hours_worked": {year: 1_500 for year in years},
+            },
+            "child": {"age": {year: 2 for year in years}},
+        },
+        "benunits": {
+            "benunit": {
+                "members": ["parent", "child"],
+                "would_claim_child_benefit": {year: True for year in years},
+            }
+        },
+        "households": {"household": {"members": ["parent", "child"]}},
+    }
+
+
+def values(simulation, variable, year=YEAR):
+    return np.asarray(simulation.calculate(variable, year)).copy()
+
+
+def add_benefit(simulation, *, read_first=True, years=(YEAR,), holder=False):
+    for year in years:
+        source = simulation if read_first else simulation.baseline
+        benefit = values(source, CHILD_BENEFIT, year)
+        if holder:
+            simulation.get_holder(CHILD_BENEFIT).set_input(
+                period(year), benefit + 1_000, simulation.branch_name
+            )
+        else:
+            simulation.set_input(CHILD_BENEFIT, year, benefit + 1_000)
+
+
+def parameter_benchmark(parameter_changes=None):
+    baseline = Simulation(situation=family())
+    weekly = baseline.tax_benefit_system.parameters.get_child(ELDEST_AMOUNT)(str(YEAR))
+    return Simulation(
+        situation=family(),
+        scenario=Scenario(
+            parameter_changes={
+                **(parameter_changes or {}),
+                ELDEST_AMOUNT: {str(YEAR): weekly + 1_000 / 52},
+            }
+        ),
+    )
+
+
+@pytest.mark.parametrize("holder", [False, True], ids=["simulation", "holder"])
+@pytest.mark.parametrize(
+    "years", [(YEAR,), (YEAR, YEAR + 1)], ids=["one-year", "two-years"]
+)
+def test_modifier_input_survives_reset(holder, years):
+    simulation = Simulation(
+        situation=family(years),
+        scenario=Scenario(
+            simulation_modifier=lambda sim: add_benefit(sim, holder=holder, years=years)
+        ),
+    )
+    assert CHILD_BENEFIT in simulation.input_variables
+    assert CHILD_BENEFIT not in simulation.baseline.input_variables
+    simulation.reset_calculations()
+    for year in years:
+        np.testing.assert_allclose(
+            values(simulation, CHILD_BENEFIT, year),
+            values(simulation.baseline, CHILD_BENEFIT, year) + 1_000,
+            atol=0.01,
+        )
+
+
+def test_modifier_inputs_exclude_calculated_intermediates():
+    def modifier(simulation):
+        values(simulation, "household_net_income")
+        add_benefit(simulation)
+
+    simulation = Simulation(
+        situation=family(), scenario=Scenario(simulation_modifier=modifier)
+    )
+    names = set(simulation.input_variables) - set(simulation.baseline.input_variables)
+    assert names == {CHILD_BENEFIT}
+    assert not names.intersection(
+        {"adjusted_net_income", "income_tax", "ni_class_1_employee"}
+    )
+    benchmark = parameter_benchmark()
+    for variable in ("household_net_income", "income_tax", "adjusted_net_income"):
+        np.testing.assert_allclose(
+            values(simulation, variable), values(benchmark, variable), atol=0.01
+        )
+
+
+@pytest.mark.parametrize("dependent", ["household_net_income", "gov_balance"])
+def test_modifier_recalculates_cached_dependents(dependent):
+    def modifier(simulation):
+        values(simulation, dependent)
+        add_benefit(simulation)
+
+    simulation = Simulation(
+        situation=family(), scenario=Scenario(simulation_modifier=modifier)
+    )
+    expected = values(parameter_benchmark(), dependent)
+    np.testing.assert_allclose(values(simulation, dependent), expected, atol=0.01)
+    simulation.reset_calculations()
+    np.testing.assert_allclose(values(simulation, dependent), expected, atol=0.01)
+
+
+def test_modifier_keeps_overwritten_input_without_mutating_baseline():
+    situation = family()
+    situation["benunits"]["benunit"][CHILD_BENEFIT] = {YEAR: 500}
+    original = copy.deepcopy(situation)
+    captured = {}
+
+    def modifier(simulation):
+        captured["names"] = simulation.baseline.input_variables
+        captured["recorded"] = list(simulation.baseline.input_variables)
+        captured["net"] = values(simulation.baseline, "household_net_income")
+        values(simulation, "household_net_income")
+        add_benefit(simulation)
+
+    simulation = Simulation(
+        situation=situation, scenario=Scenario(simulation_modifier=modifier)
+    )
+    np.testing.assert_allclose(
+        values(simulation, "household_net_income"), captured["net"] + 1_000, atol=0.01
+    )
+    simulation.reset_calculations()
+    np.testing.assert_array_equal(values(simulation, CHILD_BENEFIT), [1_500])
+    np.testing.assert_array_equal(values(simulation.baseline, CHILD_BENEFIT), [500])
+    np.testing.assert_array_equal(
+        values(simulation.baseline, "household_net_income"), captured["net"]
+    )
+    assert simulation.baseline.input_variables is captured["names"]
+    assert captured["names"] == captured["recorded"]
+    assert situation == original
+
+
+@pytest.mark.parametrize("route", ["deleted", "clone", "branch", "calculated"])
+def test_modifier_names_only_surviving_own_inputs(route):
+    def modifier(simulation):
+        if route == "deleted":
+            simulation.set_input(CHILD_BENEFIT, YEAR, [5_000])
+            simulation.delete_arrays(CHILD_BENEFIT, YEAR)
+        elif route == "clone":
+            simulation.clone().set_input(CHILD_BENEFIT, YEAR, [5_000])
+        elif route == "branch":
+            simulation.get_branch("modifier_other").set_input(
+                CHILD_BENEFIT, YEAR, [5_000]
+            )
+        else:
+            values(simulation, CHILD_BENEFIT)
+
+    simulation = Simulation(
+        situation=family(), scenario=Scenario(simulation_modifier=modifier)
+    )
+    assert simulation.input_variables == simulation.baseline.input_variables
+    simulation.reset_calculations()
+    np.testing.assert_array_equal(
+        values(simulation, CHILD_BENEFIT), values(simulation.baseline, CHILD_BENEFIT)
+    )
+
+
+@pytest.mark.parametrize(
+    "variable", ["marginal_tax_rate", "marginal_tax_rate_wrt_employer_cost"]
+)
+def test_modifier_fixed_child_benefit_addition_preserves_marginal_rates(variable):
+    simulation = Simulation(
+        situation=family(), scenario=Scenario(simulation_modifier=add_benefit)
+    )
+    np.testing.assert_allclose(
+        values(simulation, variable), values(simulation.baseline, variable), atol=1e-5
+    )
+    np.testing.assert_allclose(
+        values(simulation, variable), values(parameter_benchmark(), variable), atol=1e-5
+    )
+
+
+def test_fixed_tax_modifier_preserves_capital_gains_marginal_rate():
+    situation = family()
+    situation["people"]["parent"]["capital_gains"] = {YEAR: 10_000}
+
+    def modifier(simulation):
+        simulation.set_input("capital_gains_tax", YEAR, [5_000, 0])
+
+    simulation = Simulation(
+        situation=situation, scenario=Scenario(simulation_modifier=modifier)
+    )
+    assert values(simulation.baseline, "marginal_tax_rate_on_capital_gains")[0] > 0
+    np.testing.assert_array_equal(
+        values(simulation, "marginal_tax_rate_on_capital_gains"), [0, 0]
+    )
+    np.testing.assert_array_equal(values(simulation, "capital_gains_tax"), [5_000, 0])
+
+
+def test_modifier_dynamics_preserves_benefit_and_income_only_response():
+    simulation = Microsimulation(
+        situation=family(),
+        scenario=Scenario(
+            parameter_changes={
+                "gov.dynamic.obr_labour_supply_assumptions": {str(YEAR): True}
+            },
+            simulation_modifier=add_benefit,
+        ),
+    )
+    expected = values(simulation.baseline, CHILD_BENEFIT) + 1_000
+    progression = simulation.apply_dynamics(YEAR).progression
+    np.testing.assert_allclose(values(simulation, CHILD_BENEFIT), expected, atol=0.01)
+    np.testing.assert_allclose(
+        progression["deriv_scenario"], progression["deriv_baseline"], atol=1e-5
+    )
+    np.testing.assert_allclose(progression["substitution_response"], 0, atol=0.01)
+    assert (np.asarray(progression["income_response"]) < 0).all()
+    np.testing.assert_allclose(
+        progression["total_response"], progression["income_response"], atol=0.01
+    )
+
+
+@pytest.mark.parametrize(
+    "read_first", [False, True], ids=["baseline-set", "read-then-set"]
+)
+@pytest.mark.parametrize(
+    "first",
+    [RESPONSE, "household_net_income"],
+    ids=["response-first", "dependent-first"],
+)
+def test_modifier_response_matches_equivalent_parameter_change(read_first, first):
+    simulation = Simulation(
+        situation=family(),
+        scenario=Scenario(
+            parameter_changes=ELASTICITIES,
+            simulation_modifier=lambda sim: add_benefit(sim, read_first=read_first),
+        ),
+    )
+    benchmark = parameter_benchmark(ELASTICITIES)
+    values(simulation, first)
+    variables = (
+        CHILD_BENEFIT,
+        "household_net_income",
+        "relative_income_change",
+        "relative_wage_change",
+        "marginal_tax_rate",
+        RESPONSE,
+    )
+    expected = {variable: values(benchmark, variable) for variable in variables}
+    assert expected[RESPONSE][0] < 0
+    for variable in variables:
+        np.testing.assert_allclose(
+            values(simulation, variable),
+            expected[variable],
+            atol=0.01,
+            rtol=1e-5,
+            err_msg=variable,
+        )
+    np.testing.assert_allclose(values(simulation, "relative_wage_change"), 0, atol=1e-5)
+    np.testing.assert_allclose(values(simulation, SUBSTITUTION_TERM), 0, atol=0.01)
+    np.testing.assert_allclose(
+        values(simulation, RESPONSE), values(simulation, INCOME_TERM), atol=0.01
+    )
+    simulation.reset_calculations()
+    for variable in variables:
+        np.testing.assert_allclose(
+            values(simulation, variable),
+            expected[variable],
+            atol=0.01,
+            rtol=1e-5,
+            err_msg=variable,
+        )
+
+
+@pytest.mark.parametrize("scenario", [None, Scenario()], ids=["none", "empty"])
+def test_no_modifier_keeps_existing_behavior(scenario):
+    simulation = Simulation(situation=family(), scenario=scenario)
+    np.testing.assert_array_equal(values(simulation, RESPONSE), [0, 0])
+    np.testing.assert_array_equal(
+        values(simulation, CHILD_BENEFIT), values(simulation.baseline, CHILD_BENEFIT)
+    )
+    assert simulation.input_variables == simulation.baseline.input_variables
+
+
+def test_before_load_modifier_keeps_timing():
+    calls = []
+
+    def before(simulation):
+        calls.append(("before", getattr(simulation, "populations", None)))
+
+    def after(simulation):
+        calls.append(("after", bool(simulation.populations)))
+
+    simulation = Simulation(
+        situation=family(),
+        scenario=Scenario(simulation_modifier=before, applied_before_data_load=True)
+        + Scenario(simulation_modifier=after),
+    )
+    assert calls == [("before", None), ("after", True)]
+    assert simulation.input_variables == simulation.baseline.input_variables
+
+
+@pytest.mark.parametrize(
+    "variable", ["marginal_tax_rate", "marginal_tax_rate_wrt_employer_cost"]
+)
+def test_modifier_recalculates_cached_marginal_rate_branches(variable):
+    years = (YEAR, YEAR + 1)
+    warmed = []
+
+    def modifier(simulation):
+        for year in years:
+            warmed.append(values(simulation, variable, year))
+        add_benefit(simulation, years=years)
+
+    simulation = Simulation(
+        situation=family(years), scenario=Scenario(simulation_modifier=modifier)
+    )
+    assert len(warmed) == len(years)
+    np.testing.assert_allclose(
+        values(simulation, variable),
+        values(parameter_benchmark(), variable),
+        atol=1e-5,
+    )
+    for year in years:
+        np.testing.assert_allclose(
+            values(simulation, variable, year),
+            values(simulation.baseline, variable, year),
+            atol=1e-5,
+        )
+        np.testing.assert_allclose(
+            values(simulation, CHILD_BENEFIT, year),
+            values(simulation.baseline, CHILD_BENEFIT, year) + 1_000,
+            atol=0.01,
+        )
+
+
+def test_modifier_recalculates_warmed_capital_gains_marginal_rate():
+    situation = family()
+    situation["people"]["parent"]["capital_gains"] = {YEAR: 10_000}
+
+    def modifier(simulation):
+        assert values(simulation, "marginal_tax_rate_on_capital_gains")[0] > 0
+        # Gains-rate branches are removed by the formula after each step.
+        assert not any(name.endswith("_cg_rise") for name in simulation.branches)
+        simulation.set_input("capital_gains_tax", YEAR, [5_000, 0])
+
+    simulation = Simulation(
+        situation=situation, scenario=Scenario(simulation_modifier=modifier)
+    )
+    np.testing.assert_array_equal(
+        values(simulation, "marginal_tax_rate_on_capital_gains"), [0, 0]
+    )
+    np.testing.assert_array_equal(values(simulation, "capital_gains_tax"), [5_000, 0])
+
+
+@pytest.mark.parametrize(
+    "variable, branch_name",
+    [
+        ("marginal_tax_rate", "adult_1_pay_rise"),
+        ("marginal_tax_rate_wrt_employer_cost", "adult_1_employer_cost_mtr"),
+    ],
+)
+def test_marginal_rate_keeps_an_existing_named_branch(variable, branch_name):
+    simulation = Simulation(situation=family())
+    existing = simulation.get_branch(branch_name)
+    existing.set_input(CHILD_BENEFIT, YEAR, [5_000])
+    expected = values(parameter_benchmark(), variable)
+
+    np.testing.assert_allclose(values(simulation, variable), expected, atol=1e-5)
+
+    assert set(simulation.branches) == {branch_name}
+    assert simulation.branches[branch_name] is existing
+    np.testing.assert_array_equal(values(existing, CHILD_BENEFIT), [5_000])
+
+
+@pytest.mark.parametrize(
+    "variable", ["marginal_tax_rate", "marginal_tax_rate_wrt_employer_cost"]
+)
+def test_marginal_rate_failure_detaches_its_scratch_branch(monkeypatch, variable):
+    simulation = Simulation(situation=family())
+    assert not simulation.branches
+    set_input = Simulation.set_input
+
+    def fail_pay_rise(branch, name, *args, **kwargs):
+        if branch is not simulation and name == "employment_income":
+            raise RuntimeError("synthetic pay-rise failure")
+        return set_input(branch, name, *args, **kwargs)
+
+    monkeypatch.setattr(Simulation, "set_input", fail_pay_rise)
+    with pytest.raises(RuntimeError, match="synthetic pay-rise failure"):
+        values(simulation, variable)
+    assert not simulation.branches
diff --git a/policyengine_uk/variables/household/marginal_tax_rate.py b/policyengine_uk/variables/household/marginal_tax_rate.py
index 08d90984a7..5f7cb9d0e7 100644
--- a/policyengine_uk/variables/household/marginal_tax_rate.py
+++ b/policyengine_uk/variables/household/marginal_tax_rate.py
@@ -34,28 +34,35 @@ def formula(person, period, parameters):
         # values.
         held_responses = held_behavioural_responses(person, period)
         for adult_index in range(1, 1 + adult_count):
-            alt_simulation = simulation.get_branch(f"adult_{adult_index}_pay_rise")
-            mask = adult_index_values == adult_index
-            for variable in simulation.tax_benefit_system.variables:
-                variable_data = simulation.tax_benefit_system.variables[variable]
-                if (
-                    variable not in simulation.input_variables
-                    and not variable_data.is_input_variable()
-                ):
-                    alt_simulation.delete_arrays(variable)
-            hold_behavioural_responses(alt_simulation, period, held_responses)
-            alt_simulation.set_input(
-                "employment_income",
-                period,
-                employment_income + mask * step,
-            )
-            alt_person = alt_simulation.person
-            # Float32 rounds the higher earnings, so divide by the rise stored.
-            pay_rise = alt_person("employment_income", period) - employment_income
-            household_net_income = person.household("household_net_income", period)
-            household_net_income_higher_earnings = alt_person.household(
-                "household_net_income", period
-            )
-            increase = household_net_income_higher_earnings - household_net_income
-            mtr_values += where(mask, 1 - increase / where(mask, pay_rise, 1), 0)
+            name = f"adult_{adult_index}_pay_rise"
+            while name in simulation.branches or name == simulation.branch_name:
+                name += "_"
+            alt_simulation = simulation.get_branch(name)
+            try:
+                mask = adult_index_values == adult_index
+                for variable in simulation.tax_benefit_system.variables:
+                    variable_data = simulation.tax_benefit_system.variables[variable]
+                    if (
+                        variable not in simulation.input_variables
+                        and not variable_data.is_input_variable()
+                    ):
+                        alt_simulation.delete_arrays(variable)
+                hold_behavioural_responses(alt_simulation, period, held_responses)
+                alt_simulation.set_input(
+                    "employment_income",
+                    period,
+                    employment_income + mask * step,
+                )
+                alt_person = alt_simulation.person
+                # Float32 rounds the higher earnings, so divide by the rise stored.
+                pay_rise = alt_person("employment_income", period) - employment_income
+                household_net_income = person.household("household_net_income", period)
+                household_net_income_higher_earnings = alt_person.household(
+                    "household_net_income", period
+                )
+                increase = household_net_income_higher_earnings - household_net_income
+                mtr_values += where(mask, 1 - increase / where(mask, pay_rise, 1), 0)
+            finally:
+                if simulation.branches.get(name) is alt_simulation:
+                    del simulation.branches[name]
         return mtr_values
diff --git a/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py b/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
index a9ea335399..f71a5d0082 100644
--- a/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
+++ b/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
@@ -38,30 +38,35 @@ def formula(person, period, parameters):
         # values.
         held_responses = held_behavioural_responses(person, period)
         for adult_index in range(1, 1 + adult_count):
-            alt_simulation = simulation.get_branch(
-                f"adult_{adult_index}_employer_cost_mtr"
-            )
-            mask = adult_index_values == adult_index
-            for variable in simulation.tax_benefit_system.variables:
-                variable_data = simulation.tax_benefit_system.variables[variable]
-                if (
-                    variable not in simulation.input_variables
-                    and not variable_data.is_input_variable()
-                ):
-                    alt_simulation.delete_arrays(variable)
-            hold_behavioural_responses(alt_simulation, period, held_responses)
-            alt_simulation.set_input(
-                "employment_income",
-                period,
-                employment_income + mask * step,
-            )
-            alt_person = alt_simulation.person
-            new_net = alt_person.household("household_net_income", period)
-            new_employer_cost = alt_person("employer_cost", period)
-            net_increase = new_net - baseline_net
-            employer_cost_increase = new_employer_cost - baseline_employer_cost
-            safe_denominator = where(
-                employer_cost_increase != 0, employer_cost_increase, 1
-            )
-            mtr_values += where(mask, 1 - net_increase / safe_denominator, 0)
+            name = f"adult_{adult_index}_employer_cost_mtr"
+            while name in simulation.branches or name == simulation.branch_name:
+                name += "_"
+            alt_simulation = simulation.get_branch(name)
+            try:
+                mask = adult_index_values == adult_index
+                for variable in simulation.tax_benefit_system.variables:
+                    variable_data = simulation.tax_benefit_system.variables[variable]
+                    if (
+                        variable not in simulation.input_variables
+                        and not variable_data.is_input_variable()
+                    ):
+                        alt_simulation.delete_arrays(variable)
+                hold_behavioural_responses(alt_simulation, period, held_responses)
+                alt_simulation.set_input(
+                    "employment_income",
+                    period,
+                    employment_income + mask * step,
+                )
+                alt_person = alt_simulation.person
+                new_net = alt_person.household("household_net_income", period)
+                new_employer_cost = alt_person("employer_cost", period)
+                net_increase = new_net - baseline_net
+                employer_cost_increase = new_employer_cost - baseline_employer_cost
+                safe_denominator = where(
+                    employer_cost_increase != 0, employer_cost_increase, 1
+                )
+                mtr_values += where(mask, 1 - net_increase / safe_denominator, 0)
+            finally:
+                if simulation.branches.get(name) is alt_simulation:
+                    del simulation.branches[name]
         return mtr_values

From a703cbc2455293e320e70637dd246917cec5234e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 12:28:11 +0000
Subject: [PATCH 2/6] Reproduce direct modifier input and cache failures

---
 .../tests/test_direct_modifier_inputs.py      | 219 ++++++++++++++++++
 1 file changed, 219 insertions(+)
 create mode 100644 policyengine_uk/tests/test_direct_modifier_inputs.py

diff --git a/policyengine_uk/tests/test_direct_modifier_inputs.py b/policyengine_uk/tests/test_direct_modifier_inputs.py
new file mode 100644
index 0000000000..459f896d0c
--- /dev/null
+++ b/policyengine_uk/tests/test_direct_modifier_inputs.py
@@ -0,0 +1,219 @@
+"""Public synthetic regressions for callbacks on an already-built UK model."""
+
+import copy
+
+import numpy as np
+import pytest
+
+from policyengine_uk import Microsimulation, Simulation
+from policyengine_uk.scenarios.pip_reform import modify_simulation as phase_in_pip
+
+pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
+
+YEAR = 2027
+BENEFIT = "child_benefit"
+CLAIM = "would_claim_child_benefit"
+RESPONSE = "employment_income_behavioral_response"
+ELDEST = "gov.hmrc.child_benefit.amount.eldest"
+ELASTICITIES = {
+    "gov.simulation.labour_supply_responses.substitution_elasticity": {str(YEAR): 0.15},
+    "gov.simulation.labour_supply_responses.income_elasticity": {str(YEAR): -0.05},
+}
+RATES = (
+    "marginal_tax_rate",
+    "marginal_tax_rate_wrt_employer_cost",
+    "marginal_tax_rate_on_capital_gains",
+)
+
+
+def family(years=(YEAR,)):
+    return {
+        "people": {
+            "parent": {
+                "age": {year: 30 for year in years},
+                "gender": {year: "FEMALE" for year in years},
+                "employment_income": {year: 25_000 for year in years},
+                "hours_worked": {year: 1_500 for year in years},
+                "capital_gains": {year: 10_000 for year in years},
+            },
+            "child": {"age": {year: 2 for year in years}},
+        },
+        "benunits": {"family": {"members": ["parent", "child"]}},
+        "households": {"home": {"members": ["parent", "child"]}},
+    }
+
+
+def values(simulation, variable, year=YEAR):
+    return np.asarray(simulation.calculate(variable, year)).copy()
+
+
+def add_benefit(simulation, read_first=True, years=(YEAR,)):
+    for year in years:
+        source = simulation if read_first else simulation.baseline
+        simulation.set_input(BENEFIT, year, values(source, BENEFIT, year) + 1_000)
+    return simulation
+
+
+def benchmark(years=(YEAR,), *, claim=True, reform=None, cls=Simulation):
+    baseline = Simulation(situation=family(years))
+    amounts = {
+        str(year): (
+            baseline.tax_benefit_system.parameters.get_child(ELDEST)(str(year))
+            + 1_000 / 52
+            if claim
+            else 0
+        )
+        for year in years
+    }
+    return cls(situation=family(years), reform={**(reform or {}), ELDEST: amounts})
+
+
+@pytest.mark.parametrize("cls", [Simulation, Microsimulation])
+@pytest.mark.parametrize("read_first", [False, True], ids=["baseline-set", "read-set"])
+@pytest.mark.parametrize("first", [RESPONSE, "household_net_income"])
+def test_direct_callback_response_matches_parameter_benchmark(cls, read_first, first):
+    simulation = cls(situation=family(), reform=ELASTICITIES)
+    baseline_names = simulation.baseline.input_variables
+    original_names = list(baseline_names)
+    baseline_net = values(simulation.baseline, "household_net_income")
+    values(simulation, first)
+    # This is precisely the post-construction callback route, without Scenario.
+    add_benefit(simulation, read_first=read_first)
+    expected = benchmark(reform=ELASTICITIES, cls=cls)
+    assert values(expected, RESPONSE)[0] < 0
+    for reset in (False, True):
+        if reset:
+            simulation.reset_calculations()
+        for variable in (
+            BENEFIT, "household_net_income", "relative_income_change",
+            "relative_wage_change", "marginal_tax_rate", RESPONSE,
+            "income_elasticity_lsr", "substitution_elasticity_lsr",
+        ):
+            np.testing.assert_allclose(
+                values(simulation, variable), values(expected, variable),
+                atol=0.01, rtol=1e-5, err_msg=variable,
+            )
+        np.testing.assert_allclose(values(simulation, "relative_wage_change"), 0, atol=1e-5)
+        np.testing.assert_allclose(values(simulation, "substitution_elasticity_lsr"), 0, atol=0.01)
+    assert simulation.baseline.input_variables is baseline_names
+    assert baseline_names == original_names
+    np.testing.assert_array_equal(values(simulation.baseline, "household_net_income"), baseline_net)
+
+
+@pytest.mark.parametrize("claim", [True, False], ids=["formula-override", "input-only"])
+def test_direct_writes_survive_reset_marginal_rates_and_export(claim):
+    years = (YEAR, YEAR + 1)
+    situation = family(years)
+    original = copy.deepcopy(situation)
+    simulation = Simulation(situation=situation)
+    variable = BENEFIT if claim else CLAIM
+    assert variable not in simulation.input_variables
+    values(simulation, "household_net_income")
+    if claim:
+        add_benefit(simulation, years=years)
+    else:
+        for year in years:
+            simulation.set_input(CLAIM, year, [False])
+    expected = benchmark(years, claim=claim)
+    for year in years:
+        for output in (*RATES, BENEFIT, "household_net_income"):
+            np.testing.assert_allclose(
+                values(simulation, output, year), values(expected, output, year),
+                atol=0.01, rtol=1e-5, err_msg=output,
+            )
+    simulation.reset_calculations()
+    for year in years:
+        np.testing.assert_allclose(values(simulation, BENEFIT, year), values(expected, BENEFIT, year), atol=0.01)
+    frame = simulation.to_input_dataframe()
+    assert all(f"{variable}__{year}" in frame for year in years)
+    assert not any(column.startswith("household_net_income__") for column in frame)
+    rebuilt = Simulation(dataset=frame)
+    for year in years:
+        np.testing.assert_allclose(values(rebuilt, BENEFIT, year), values(expected, BENEFIT, year), atol=0.01)
+    assert variable not in simulation.baseline.input_variables
+    assert situation == original
+
+
+@pytest.mark.parametrize("claim", [True, False], ids=["formula-override", "input-only"])
+def test_direct_writes_survive_obr_dynamics(claim):
+    reform = {"gov.dynamic.obr_labour_supply_assumptions": {str(YEAR): True}}
+    simulation = Microsimulation(situation=family(), reform=reform)
+    values(simulation, "household_net_income")
+    if claim:
+        add_benefit(simulation)
+    else:
+        simulation.set_input(CLAIM, YEAR, [False])
+    expected = benchmark(claim=claim, reform=reform, cls=Microsimulation)
+    result = simulation.apply_dynamics(YEAR).progression
+    comparison = expected.apply_dynamics(YEAR).progression
+    np.testing.assert_allclose(values(simulation, BENEFIT), values(expected, BENEFIT), atol=0.01)
+    for variable in ("deriv_scenario", "substitution_response", "income_response", "total_response"):
+        np.testing.assert_allclose(result[variable], comparison[variable], atol=0.01, rtol=1e-5)
+    assert np.any(np.asarray(comparison["income_response"]) != 0)
+
+
+def test_direct_fixed_capital_gains_tax_has_zero_marginal_rate():
+    simulation = Simulation(situation=family())
+    assert values(simulation, RATES[2])[0] > 0
+    simulation.set_input("capital_gains_tax", YEAR, [5_000, 0])
+    np.testing.assert_array_equal(values(simulation, RATES[2]), [0, 0])
+    simulation.reset_calculations()
+    np.testing.assert_array_equal(values(simulation, "capital_gains_tax"), [5_000, 0])
+
+
+@pytest.mark.parametrize("route", ["calculated", "deleted", "clone", "other-branch"])
+def test_direct_retention_excludes_calculated_deleted_and_foreign_inputs(route):
+    simulation = Simulation(situation=family())
+    original = list(simulation.input_variables)
+    if route == "calculated":
+        values(simulation, "household_net_income")
+    elif route == "deleted":
+        simulation.set_input(BENEFIT, YEAR, [0])
+        simulation.delete_arrays(BENEFIT, YEAR)
+    elif route == "clone":
+        simulation.clone().set_input(BENEFIT, YEAR, [0])
+    else:
+        simulation.get_branch("owned-by-user").set_input(BENEFIT, YEAR, [0])
+    simulation.reset_calculations()
+    assert simulation.input_variables == original
+    np.testing.assert_array_equal(values(simulation, BENEFIT), values(simulation.baseline, BENEFIT))
+
+
+def test_direct_zero_overwrite_and_inherited_branch_are_isolated():
+    situation = family()
+    situation["benunits"]["family"][BENEFIT] = {YEAR: 500}
+    simulation = Simulation(situation=situation)
+    simulation.set_input(BENEFIT, YEAR, [0])
+    branch = simulation.get_branch("owned-by-user")
+    branch.reset_calculations()
+    np.testing.assert_array_equal(values(branch, BENEFIT), [0])
+    simulation.reset_calculations()
+    np.testing.assert_array_equal(values(simulation, BENEFIT), [0])
+    np.testing.assert_array_equal(values(simulation.baseline, BENEFIT), [500])
+    assert simulation.branches["owned-by-user"] is branch
+
+
+def test_shipped_pip_callback_retains_its_explicit_periods():
+    years = tuple(range(2025, 2030))
+    people = {
+        f"person{i}": {
+            "age": {year: 30 for year in years},
+            "pip_dl": {year: 100 for year in years},
+            "pip_m": {year: 0 for year in years},
+        }
+        for i in range(8)
+    }
+    simulation = Simulation(situation={
+        "people": people,
+        "benunits": {name: {"members": [name]} for name in people},
+        "households": {name: {"members": [name]} for name in people},
+    })
+    phase_in_pip(simulation)
+    expected = {year: values(simulation, "pip", year) for year in years}
+    assert np.any(expected[2029] == 0) and np.any(expected[2029] > 0)
+    simulation.reset_calculations()
+    for year in years:
+        np.testing.assert_array_equal(values(simulation, "pip", year), expected[year])
+        np.testing.assert_array_equal(values(simulation.baseline, "pip", year), np.full(8, 100))
+    frame = simulation.to_input_dataframe()
+    assert all(f"pip__{year}" in frame for year in years)

From 12bd4f3c219541103281f7d1515e2d2720d6ba69 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 13:02:26 +0000
Subject: [PATCH 3/6] Preserve direct modifier inputs through resets and
 measurements

---
 ...fix-scenario-modifier-inputs-2666.fixed.md |  2 +-
 docs/book/usage/scenarios.md                  | 23 ++++-
 policyengine_uk/simulation.py                 | 63 ++++++++-----
 .../tests/test_direct_modifier_inputs.py      | 89 ++++++++++++++-----
 .../tests/test_reform_inputs_survive_reset.py | 53 ++++++-----
 policyengine_uk/utils/measurement_branches.py | 24 ++++-
 .../marginal_tax_rate_on_capital_gains.py     |  7 +-
 .../variables/household/marginal_tax_rate.py  |  4 +-
 .../marginal_tax_rate_wrt_employer_cost.py    |  4 +-
 9 files changed, 189 insertions(+), 80 deletions(-)

diff --git a/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
index 12873d9ef7..d0cf179632 100644
--- a/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
+++ b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
@@ -1 +1 @@
-Keep inputs stored by construction-time scenario modifiers through resets and marginal-rate calculations, and clear their stale calculated values so labour supply responses reflect the completed scenario.
+Keep inputs stored by scenario modifiers and direct callbacks through resets and all marginal-rate calculations, including input-only variables absent at construction. Refresh dependent values after direct input edits so labour supply responses use the completed inputs, while preserving calculation-time framework writes and independent baselines.
diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md
index 388bfadc42..44fbbc32d4 100644
--- a/docs/book/usage/scenarios.md
+++ b/docs/book/usage/scenarios.md
@@ -175,9 +175,26 @@ An input remains fixed when a marginal-rate calculation changes earnings. A
 modifier that adds a fixed amount to Child Benefit therefore keeps that amount
 through the earnings step. For a policy that should vary with earnings, use
 parameter changes or a formula reform; setting its calculated value once cannot
-represent that dependence. This cleanup happens when the constructor runs the
-modifier. Calling a modifier directly, or using `Scenario.apply` on an existing
-simulation, retains those operations' existing behavior.
+represent that dependence.
+
+The same input retention applies when a callback calls `set_input` directly on
+an already-built simulation, including through `Scenario.apply`. Resets and all
+three marginal-rate calculations retain surviving inputs visible to that
+simulation, including inputs absent at construction. A top-level `set_input`
+also refreshes dependent calculations before the next read, so a callback can
+read Child Benefit, add an amount and then obtain fresh income and response
+values. Inputs stored only on another clone or branch do not become the
+parent's inputs. Formula and framework writes during a calculation keep their
+existing calculation lifecycle.
+
+Direct edits of the parameter tree require a separate array-cache boundary.
+After calling a parameter's `update` method on a built simulation, call
+`simulation.drop_computed_arrays()` before reading its results. In a wrapper
+that runs a policy callback, late dynamic parameter updates and a dynamic
+callback, clear computed arrays after each operation in that order. This
+preserves the input values and lets the next operation see the updated policy.
+A wrapper must also use a country-model release containing the native input
+retention fix; array cleanup alone does not add it to an older release.
 
 Modifiers with `applied_before_data_load=True` run before the simulation has its
 data. They can change parameters or model structure at that stage; modifiers
diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
index efab5cd213..0c9526ad78 100644
--- a/policyengine_uk/simulation.py
+++ b/policyengine_uk/simulation.py
@@ -2,7 +2,7 @@
 import os
 import copy
 import warnings
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from importlib.metadata import version
 from typing import Dict, Optional, Union, Type, List
 
@@ -58,7 +58,10 @@
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.data_source import built_from_data
 from policyengine_uk.utils.input_dataframe_export import input_dataframe
-from policyengine_uk.utils.measurement_branches import drop_calculated_values
+from policyengine_uk.utils.measurement_branches import (
+    drop_calculated_values,
+    retained_input_variables,
+)
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
     copy_model_values,
@@ -286,6 +289,7 @@ def __init__(
         self._cgt_input_provenance = {}
         self._cgt_defaulted_periods = set()
         self._cgt_calculation_depth = 0
+        self._external_input_change_pending = False
         self._cgt_operation_inputs = set()
         self._cgt_warned_unused = False
         self._cgt_warning_enabled = False
@@ -461,34 +465,31 @@ def _refresh_capital_gains_projection(self) -> None:
         self.tax_benefit_system.reset_parameter_caches()
 
     @contextmanager
-    def _preserve_stored_inputs(self):
-        """Name inputs an operation stores and retains, for later resets."""
+    def _record_stored_inputs(self):
+        """Record actual stores, including overwrites and custom input routes."""
         logs = _input_key_logs.setdefault(id(self), [])
         stored = set()
         logs.append(stored)
         self._user_input_keys = getattr(self, "_user_input_keys", None) or set()
         try:
-            yield
+            yield stored
         finally:
             logs.pop()
             if not logs:
                 del _input_key_logs[id(self)]
             # A plain set again, unless an enclosing operation is still logging.
             self._user_input_keys = set(self._user_input_keys)
-        # input_variables is taken when the simulation is built.
-        # reset_calculations deletes the values of every variable it does not
-        # name (as core's derivative did before policyengine-core 3.33.0, from
-        # which it keeps every input), and the marginal rate formulas those of
-        # every such variable that has a formula. A reform or scenario modifier
-        # applied afterwards can store inputs: new ones, or ones that replace
-        # an input set since the build. Name every variable with an input it
-        # stored that this simulation still holds, under a branch name it
-        # reads and a reset deletes, or a reset would delete what the reform
-        # stored and recalculate the formula. input_variables names whole
-        # variables, so a reset then keeps every input of the variable, any
-        # set since the build included. An operation that raises names nothing
-        # it stored itself (a reform it applied first has named its own). Assign
-        # a new list: a clone or branch shares its parent's.
+
+    @contextmanager
+    def _preserve_stored_inputs(self):
+        """Name inputs an operation stores and retains, for later resets."""
+        with self._record_stored_inputs() as stored:
+            yield
+        # Preserve the existing construction/reform input-name contract.
+        # Later direct stores are also retained by retained_input_variables;
+        # calculated values do not become inputs. An operation that raises
+        # names nothing it stored itself (a reform it applied first has named
+        # its own). Assign a new list: a clone or branch shares its parent's.
         input_variables = getattr(self, "input_variables", None)
         if input_variables is not None:
             branches = set(self._get_visible_branch_names())
@@ -578,10 +579,25 @@ def set_input(self, variable_name: str, period: Period, value) -> None:
         # stores a value for, and only those. set_model_value records its
         # value as the model's after this call.
         records = model_value_input_records(self, variable_name)
+        # This depth covers every calculate call, including nongains
+        # variables. Formula/framework writes must not invalidate a calculation
+        # still running; direct edits of a built simulation must refresh its
+        # dependent values before the next outermost calculation.
+        external = (
+            getattr(self, "input_variables", None) is not None
+            and self._cgt_calculation_depth == 0
+        )
+        stored = set()
         try:
-            super().set_input(variable_name, period, value)
+            recording = (
+                self._record_stored_inputs() if external else nullcontext(stored)
+            )
+            with recording as stored:
+                super().set_input(variable_name, period, value)
         finally:
             forget_replaced_model_values(self, variable_name, records)
+            if external and stored:
+                self._external_input_change_pending = True
         # Inputs do not carry back, so give the simulation its identifiers
         # for the year the input starts in if it holds them only for later
         # years. See ``policyengine_uk.data.input_dataframe.hold_identifiers_in``.
@@ -851,8 +867,10 @@ def _dataset_cache_key(self, source: str):
         )
 
     def reset_calculations(self):
+        retained = retained_input_variables(self)
+        self._external_input_change_pending = False
         for variable in self.tax_benefit_system.variables:
-            if variable not in self.input_variables:
+            if variable not in retained:
                 self.delete_arrays(variable)
 
     def get_known_variables(self):
@@ -1419,6 +1437,9 @@ def calculate(
         map_to: str = None,
         decode_enums: bool = False,
     ):
+        if self._cgt_calculation_depth == 0 and self._external_input_change_pending:
+            self._external_input_change_pending = False
+            drop_calculated_values(self)
         if variable_name is None:
             return self.calculate_all()
         tracer: SimpleTracer = self.tracer
diff --git a/policyengine_uk/tests/test_direct_modifier_inputs.py b/policyengine_uk/tests/test_direct_modifier_inputs.py
index 459f896d0c..b9f9d70170 100644
--- a/policyengine_uk/tests/test_direct_modifier_inputs.py
+++ b/policyengine_uk/tests/test_direct_modifier_inputs.py
@@ -85,25 +85,46 @@ def test_direct_callback_response_matches_parameter_benchmark(cls, read_first, f
         if reset:
             simulation.reset_calculations()
         for variable in (
-            BENEFIT, "household_net_income", "relative_income_change",
-            "relative_wage_change", "marginal_tax_rate", RESPONSE,
-            "income_elasticity_lsr", "substitution_elasticity_lsr",
+            BENEFIT,
+            "household_net_income",
+            "relative_income_change",
+            "relative_wage_change",
+            "marginal_tax_rate",
+            RESPONSE,
+            "income_elasticity_lsr",
+            "substitution_elasticity_lsr",
         ):
             np.testing.assert_allclose(
-                values(simulation, variable), values(expected, variable),
-                atol=0.01, rtol=1e-5, err_msg=variable,
+                values(simulation, variable),
+                values(expected, variable),
+                atol=0.01,
+                rtol=1e-5,
+                err_msg=variable,
             )
-        np.testing.assert_allclose(values(simulation, "relative_wage_change"), 0, atol=1e-5)
-        np.testing.assert_allclose(values(simulation, "substitution_elasticity_lsr"), 0, atol=0.01)
+        np.testing.assert_allclose(
+            values(simulation, "relative_wage_change"), 0, atol=1e-5
+        )
+        np.testing.assert_allclose(
+            values(simulation, "substitution_elasticity_lsr"), 0, atol=0.01
+        )
     assert simulation.baseline.input_variables is baseline_names
     assert baseline_names == original_names
-    np.testing.assert_array_equal(values(simulation.baseline, "household_net_income"), baseline_net)
+    np.testing.assert_array_equal(
+        values(simulation.baseline, "household_net_income"), baseline_net
+    )
 
 
 @pytest.mark.parametrize("claim", [True, False], ids=["formula-override", "input-only"])
 def test_direct_writes_survive_reset_marginal_rates_and_export(claim):
     years = (YEAR, YEAR + 1)
     situation = family(years)
+    # Export replay requires explicit identifiers, as dataset inputs do.
+    for identifier, person in enumerate(situation["people"].values(), start=1):
+        person["person_id"] = {year: identifier for year in years}
+        person["person_benunit_id"] = {year: 1 for year in years}
+        person["person_household_id"] = {year: 1 for year in years}
+    situation["benunits"]["family"]["benunit_id"] = {year: 1 for year in years}
+    situation["households"]["home"]["household_id"] = {year: 1 for year in years}
     original = copy.deepcopy(situation)
     simulation = Simulation(situation=situation)
     variable = BENEFIT if claim else CLAIM
@@ -118,18 +139,27 @@ def test_direct_writes_survive_reset_marginal_rates_and_export(claim):
     for year in years:
         for output in (*RATES, BENEFIT, "household_net_income"):
             np.testing.assert_allclose(
-                values(simulation, output, year), values(expected, output, year),
-                atol=0.01, rtol=1e-5, err_msg=output,
+                values(simulation, output, year),
+                values(expected, output, year),
+                atol=0.01,
+                rtol=1e-5,
+                err_msg=output,
             )
     simulation.reset_calculations()
     for year in years:
-        np.testing.assert_allclose(values(simulation, BENEFIT, year), values(expected, BENEFIT, year), atol=0.01)
+        np.testing.assert_allclose(
+            values(simulation, BENEFIT, year),
+            values(expected, BENEFIT, year),
+            atol=0.01,
+        )
     frame = simulation.to_input_dataframe()
     assert all(f"{variable}__{year}" in frame for year in years)
     assert not any(column.startswith("household_net_income__") for column in frame)
     rebuilt = Simulation(dataset=frame)
     for year in years:
-        np.testing.assert_allclose(values(rebuilt, BENEFIT, year), values(expected, BENEFIT, year), atol=0.01)
+        np.testing.assert_allclose(
+            values(rebuilt, BENEFIT, year), values(expected, BENEFIT, year), atol=0.01
+        )
     assert variable not in simulation.baseline.input_variables
     assert situation == original
 
@@ -146,9 +176,18 @@ def test_direct_writes_survive_obr_dynamics(claim):
     expected = benchmark(claim=claim, reform=reform, cls=Microsimulation)
     result = simulation.apply_dynamics(YEAR).progression
     comparison = expected.apply_dynamics(YEAR).progression
-    np.testing.assert_allclose(values(simulation, BENEFIT), values(expected, BENEFIT), atol=0.01)
-    for variable in ("deriv_scenario", "substitution_response", "income_response", "total_response"):
-        np.testing.assert_allclose(result[variable], comparison[variable], atol=0.01, rtol=1e-5)
+    np.testing.assert_allclose(
+        values(simulation, BENEFIT), values(expected, BENEFIT), atol=0.01
+    )
+    for variable in (
+        "deriv_scenario",
+        "substitution_response",
+        "income_response",
+        "total_response",
+    ):
+        np.testing.assert_allclose(
+            result[variable], comparison[variable], atol=0.01, rtol=1e-5
+        )
     assert np.any(np.asarray(comparison["income_response"]) != 0)
 
 
@@ -176,7 +215,9 @@ def test_direct_retention_excludes_calculated_deleted_and_foreign_inputs(route):
         simulation.get_branch("owned-by-user").set_input(BENEFIT, YEAR, [0])
     simulation.reset_calculations()
     assert simulation.input_variables == original
-    np.testing.assert_array_equal(values(simulation, BENEFIT), values(simulation.baseline, BENEFIT))
+    np.testing.assert_array_equal(
+        values(simulation, BENEFIT), values(simulation.baseline, BENEFIT)
+    )
 
 
 def test_direct_zero_overwrite_and_inherited_branch_are_isolated():
@@ -203,17 +244,21 @@ def test_shipped_pip_callback_retains_its_explicit_periods():
         }
         for i in range(8)
     }
-    simulation = Simulation(situation={
-        "people": people,
-        "benunits": {name: {"members": [name]} for name in people},
-        "households": {name: {"members": [name]} for name in people},
-    })
+    simulation = Simulation(
+        situation={
+            "people": people,
+            "benunits": {name: {"members": [name]} for name in people},
+            "households": {name: {"members": [name]} for name in people},
+        }
+    )
     phase_in_pip(simulation)
     expected = {year: values(simulation, "pip", year) for year in years}
     assert np.any(expected[2029] == 0) and np.any(expected[2029] > 0)
     simulation.reset_calculations()
     for year in years:
         np.testing.assert_array_equal(values(simulation, "pip", year), expected[year])
-        np.testing.assert_array_equal(values(simulation.baseline, "pip", year), np.full(8, 100))
+        np.testing.assert_array_equal(
+            values(simulation.baseline, "pip", year), np.full(8, 100)
+        )
     frame = simulation.to_input_dataframe()
     assert all(f"pip__{year}" in frame for year in years)
diff --git a/policyengine_uk/tests/test_reform_inputs_survive_reset.py b/policyengine_uk/tests/test_reform_inputs_survive_reset.py
index 7daf7f62a8..b1d1497892 100644
--- a/policyengine_uk/tests/test_reform_inputs_survive_reset.py
+++ b/policyengine_uk/tests/test_reform_inputs_survive_reset.py
@@ -1,14 +1,12 @@
 """A reform applied to a built simulation keeps the inputs it stores.
 
-input_variables is taken when a simulation is built, and reset_calculations
-deletes the values of every variable it does not name (as derivative did
-before policyengine-core 3.33.0, from which it keeps every input). So
-Simulation.apply_reform names each variable the reform stored an input for
-that the simulation still holds afterwards, however the reform stored it, and
-no other: an input set since the build that the reform did not store is still
-deleted by a reset, as before, unless the reform names its variable.
-input_variables names whole variables, so a reset keeps every input of a named
-variable. See _InputKeyLog in policyengine_uk/simulation.py.
+input_variables starts with the build's inputs. Simulation.apply_reform names
+each variable the reform stored an input for that the simulation still holds,
+however the reform stored it, and no other. Its own-write naming contract is
+separate from reset_calculations, which also preserves surviving explicit
+inputs supplied directly since construction. Deletion removes those inputs;
+calculated values do not become supplied inputs. See _InputKeyLog in
+policyengine_uk/simulation.py.
 
 Each simulation is one adult, the claimant, with £3,000 of income-related ESA
 entered for 2026, so their claimant-or-partner award that year is £3,000
@@ -211,13 +209,23 @@ def sets_a_year_whose_months_a_helper_keeps(simulation):
 @pytest.mark.parametrize("case", NOT_STORING)
 def test_a_reform_names_no_input_it_did_not_store(case):
     # The reform tries to set the variable but leaves no input of its own,
-    # while one set since the build survives. That one is not the reform's,
-    # so it is not named, and a reset deletes it as it would have anyway.
+    # while one supplied directly since construction survives. That input is
+    # not named by the reform, but reset still preserves the explicit value.
     simulation = built()
     recorded = list(simulation.input_variables)
     variable, year, calculated = NOT_STORING[case](simulation)
     assert simulation.input_variables == recorded
     simulation.reset_calculations()
+    supplied = {
+        "stores and deletes another year": 100,
+        "sets a year after the variable's end": 100,
+        "sets a year whose months a helper keeps": True,
+    }[case]
+    assert simulation.calculate(variable, year).tolist() == [supplied]
+    # Differential control: deleting that actual input restores the original
+    # formula/default. Retention does not freeze calculated values as inputs.
+    simulation.delete_arrays(variable, year)
+    simulation.reset_calculations()
     assert simulation.calculate(variable, year).tolist() == [calculated]
 
 
@@ -321,10 +329,9 @@ def apply(reform):
 
 # Property: drawn inputs set since the build, then a reform that sets and
 # deletes inputs by any route. The reform names exactly the variables it
-# leaves an input of its own for. A reset keeps every input of those
-# variables (its own, and any set since the build that it left in place),
-# and every variable it does not name recalculates as if nothing had been
-# set since the build.
+# leaves an input of its own for. Reset independently keeps every surviving
+# supplied cell, including earlier direct stores the reform did not write.
+# A genuinely unsupplied cell recalculates as it does in the pristine model.
 VARIABLES = [SCOPED, "claimant_or_partner_jsa_income"]
 YEARS = [AWARD_YEAR - 1, AWARD_YEAR, AWARD_YEAR + 1]
 ROUTES = ["simulation", "holder", "thread", "delete"]
@@ -400,13 +407,13 @@ def apply(reform):
     assert names == recorded
     simulation.reset_calculations()
     for (variable, year), value in left.items():
-        if variable in named:
-            assert simulation.calculate(variable, year).tolist() == [value]
+        assert simulation.calculate(variable, year).tolist() == [value]
     pristine = parent.clone()
-    for variable in set(VARIABLES) - named:
+    for variable in VARIABLES:
         for year in YEARS:
-            np.testing.assert_array_equal(
-                simulation.calculate(variable, year),
-                pristine.calculate(variable, year),
-                err_msg=f"{variable} {year}",
-            )
+            if (variable, year) not in left:
+                np.testing.assert_array_equal(
+                    simulation.calculate(variable, year),
+                    pristine.calculate(variable, year),
+                    err_msg=f"{variable} {year}",
+                )
diff --git a/policyengine_uk/utils/measurement_branches.py b/policyengine_uk/utils/measurement_branches.py
index bd6c901cd1..672eca5eac 100644
--- a/policyengine_uk/utils/measurement_branches.py
+++ b/policyengine_uk/utils/measurement_branches.py
@@ -23,6 +23,25 @@
 )
 
 
+def retained_input_variables(simulation) -> set:
+    """Names retained by resets and measurements, including later inputs.
+
+    A key can outlive its array or belong to a branch this simulation cannot
+    read. Check the supplied storage record as well as branch visibility, so
+    a calculated or deleted value does not become a fixed input.
+    """
+    retained = set(getattr(simulation, "input_variables", None) or ())
+    visible = set(simulation._get_visible_branch_names())
+    for name, branch, period in tuple(
+        getattr(simulation, "_user_input_keys", None) or ()
+    ):
+        if branch in visible and simulation.get_holder(name)._stores_user_input(
+            period, branch
+        ):
+            retained.add(name)
+    return retained
+
+
 def drop_calculated_values(simulation) -> None:
     """Delete the values a simulation has calculated, keeping its inputs.
 
@@ -31,10 +50,7 @@ def drop_calculated_values(simulation) -> None:
     simulation modifier does to formula variables. Only variables the
     simulation holds values for are visited.
     """
-    supplied = set(getattr(simulation, "input_variables", None) or ())
-    supplied.update(
-        name for name, _, _ in getattr(simulation, "_user_input_keys", None) or ()
-    )
+    supplied = retained_input_variables(simulation)
     variables = simulation.tax_benefit_system.variables
     held = {
         name
diff --git a/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py b/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
index c12c8c9e46..7b6948b849 100644
--- a/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
+++ b/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
@@ -3,6 +3,7 @@
 from policyengine_uk.utils.measurement_branches import (
     held_behavioural_responses,
     hold_behavioural_responses,
+    retained_input_variables,
 )
 from policyengine_uk.utils.marginal_rates import marginal_rate_step
 
@@ -46,12 +47,10 @@ def formula(person, period, parameters):
         for adult_index in [1, 2]:
             alt_simulation = simulation.get_branch(f"adult_{adult_index}_cg_rise")
             mask = adult_index_values == adult_index
+            retained = retained_input_variables(alt_simulation)
             for variable in simulation.tax_benefit_system.variables:
                 variable_data = simulation.tax_benefit_system.variables[variable]
-                if (
-                    variable not in simulation.input_variables
-                    and not variable_data.is_input_variable()
-                ):
+                if variable not in retained and not variable_data.is_input_variable():
                     alt_simulation.delete_arrays(variable)
             hold_behavioural_responses(alt_simulation, period, held_responses)
             alt_simulation.set_input("capital_gains", period, gains + mask * step)
diff --git a/policyengine_uk/variables/household/marginal_tax_rate.py b/policyengine_uk/variables/household/marginal_tax_rate.py
index 5f7cb9d0e7..50f8c5208b 100644
--- a/policyengine_uk/variables/household/marginal_tax_rate.py
+++ b/policyengine_uk/variables/household/marginal_tax_rate.py
@@ -4,6 +4,7 @@
 from policyengine_uk.utils.measurement_branches import (
     held_behavioural_responses,
     hold_behavioural_responses,
+    retained_input_variables,
 )
 
 
@@ -40,10 +41,11 @@ def formula(person, period, parameters):
             alt_simulation = simulation.get_branch(name)
             try:
                 mask = adult_index_values == adult_index
+                retained = retained_input_variables(alt_simulation)
                 for variable in simulation.tax_benefit_system.variables:
                     variable_data = simulation.tax_benefit_system.variables[variable]
                     if (
-                        variable not in simulation.input_variables
+                        variable not in retained
                         and not variable_data.is_input_variable()
                     ):
                         alt_simulation.delete_arrays(variable)
diff --git a/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py b/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
index f71a5d0082..5be93cbc6a 100644
--- a/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
+++ b/policyengine_uk/variables/household/marginal_tax_rate_wrt_employer_cost.py
@@ -4,6 +4,7 @@
 from policyengine_uk.utils.measurement_branches import (
     held_behavioural_responses,
     hold_behavioural_responses,
+    retained_input_variables,
 )
 
 
@@ -44,10 +45,11 @@ def formula(person, period, parameters):
             alt_simulation = simulation.get_branch(name)
             try:
                 mask = adult_index_values == adult_index
+                retained = retained_input_variables(alt_simulation)
                 for variable in simulation.tax_benefit_system.variables:
                     variable_data = simulation.tax_benefit_system.variables[variable]
                     if (
-                        variable not in simulation.input_variables
+                        variable not in retained
                         and not variable_data.is_input_variable()
                     ):
                         alt_simulation.delete_arrays(variable)

From ace3990bdba6f2ab09bba6d6e3fe5b8b3d5ac3fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 15:54:45 +0000
Subject: [PATCH 4/6] Keep formula-owned branch caches during modifier refresh

---
 policyengine_uk/simulation.py                 | 16 ++--
 ...test_cached_arrays_not_written_in_place.py | 22 +++++-
 .../test_capital_gains_response_default.py    | 16 +++-
 .../tests/test_direct_modifier_inputs.py      | 77 +++++++++++++++++++
 .../tests/test_property_income_properties.py  |  7 ++
 5 files changed, 128 insertions(+), 10 deletions(-)

diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
index 0c9526ad78..735ae8032e 100644
--- a/policyengine_uk/simulation.py
+++ b/policyengine_uk/simulation.py
@@ -580,12 +580,18 @@ def set_input(self, variable_name: str, period: Period, value) -> None:
         # value as the model's after this call.
         records = model_value_input_records(self, variable_name)
         # This depth covers every calculate call, including nongains
-        # variables. Formula/framework writes must not invalidate a calculation
-        # still running; direct edits of a built simulation must refresh its
-        # dependent values before the next outermost calculation.
+        # variables. A formula can write to a new branch whose own depth is
+        # zero while an ancestor is calculating. Core already invalidates
+        # that branch's affected dependencies; keep its unrelated caches.
+        # Direct edits with every ancestor idle must refresh dependent values
+        # before the next outermost calculation.
+        calculation = self
+        while calculation is not None and not getattr(
+            calculation, "_cgt_calculation_depth", 0
+        ):
+            calculation = getattr(calculation, "parent_branch", None)
         external = (
-            getattr(self, "input_variables", None) is not None
-            and self._cgt_calculation_depth == 0
+            getattr(self, "input_variables", None) is not None and calculation is None
         )
         stored = set()
         try:
diff --git a/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py b/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py
index 1d688833ce..0b2ea89e93 100644
--- a/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py
+++ b/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py
@@ -550,17 +550,33 @@ def _earning_couple_with_gains() -> dict:
         "capital_gains_default_response",
     ],
 )
-def test_branching_runs_with_read_only_cache(read_only_cache, case):
+def test_branching_runs_with_read_only_cache(read_only_cache, case, monkeypatch):
     from policyengine_uk import Simulation
     from policyengine_uk.model_api import Scenario
 
     if case == "marginal_rates":
         sim = Simulation(situation=_earning_couple_with_gains())
+        owned = []
+        original_get_branch = Simulation.get_branch
+
+        def capture(owner, name="branch", *args, **kwargs):
+            branch = original_get_branch(owner, name, *args, **kwargs)
+            if owner is sim and name.startswith(
+                ("adult_1_pay_rise", "adult_2_pay_rise")
+            ):
+                owned.append((name, branch))
+            return branch
+
+        monkeypatch.setattr(Simulation, "get_branch", capture)
         for year in YEARS:
             sim.calculate("marginal_tax_rate", year)
             sim.calculate("marginal_tax_rate_on_capital_gains", year)
-        # The branches ran: both earners face a rate, and so do the gains.
-        assert {"adult_1_pay_rise", "adult_2_pay_rise"} <= set(sim.branches)
+        # Both scratch branches ran, and each registration was detached.
+        assert {"adult_1_pay_rise", "adult_2_pay_rise"} <= {
+            name for name, branch in owned
+        }
+        assert all(sim.branches.get(name) is not branch for name, branch in owned)
+        # Both earners face a rate, and so do the gains.
         assert (sim.calculate("marginal_tax_rate", 2026)[:2] > 0).all()
         assert sim.calculate("marginal_tax_rate_on_capital_gains", 2026)[0] > 0
         return
diff --git a/policyengine_uk/tests/test_capital_gains_response_default.py b/policyengine_uk/tests/test_capital_gains_response_default.py
index ae2ffbd62f..a3224d8bd6 100644
--- a/policyengine_uk/tests/test_capital_gains_response_default.py
+++ b/policyengine_uk/tests/test_capital_gains_response_default.py
@@ -463,7 +463,7 @@ def test_the_rate_on_gains_is_the_share_net_income_does_not_keep(earnings, gains
 # Marginal tax rates hold realisations fixed.
 
 
-def test_pay_rise_branches_hold_the_realisation_response():
+def test_pay_rise_branches_hold_the_realisation_response(monkeypatch):
     """A pay rise is not a reform for the response to answer.
 
     The marginal tax rate's pay-rise branch measured the reform's rate on
@@ -479,15 +479,27 @@ def test_pay_rise_branches_hold_the_realisation_response():
         situation=single(40_000, realised),
         reform={**EQUALISED_RATES, **elasticities(elasticity=0)},
     )
+    owned = []
+    original_get_branch = Simulation.get_branch
+
+    def capture(owner, name="branch", *args, **kwargs):
+        branch = original_get_branch(owner, name, *args, **kwargs)
+        if owner is reform and name.startswith("adult_") and "pay_rise" in name:
+            owned.append((name, branch))
+        return branch
+
+    monkeypatch.setattr(Simulation, "get_branch", capture)
 
     rate = values(reform, "marginal_tax_rate")[0]
     assert realised < 12_770
     assert rate == pytest.approx(values(static, "marginal_tax_rate")[0], abs=1e-6)
     assert 0 < rate < 1
-    branch = reform.branches["adult_1_pay_rise"]
+    assert owned
+    branch = owned[0][1]
     assert list(values(branch, "capital_gains_behavioural_response")) == list(
         values(reform, "capital_gains_behavioural_response")
     )
+    assert all(reform.branches.get(name) is not branch for name, branch in owned)
 
 
 def test_employer_cost_branches_hold_the_realisation_response():
diff --git a/policyengine_uk/tests/test_direct_modifier_inputs.py b/policyengine_uk/tests/test_direct_modifier_inputs.py
index b9f9d70170..0785f8cfc8 100644
--- a/policyengine_uk/tests/test_direct_modifier_inputs.py
+++ b/policyengine_uk/tests/test_direct_modifier_inputs.py
@@ -4,8 +4,11 @@
 
 import numpy as np
 import pytest
+from policyengine_core.periods import YEAR as YEAR_PERIOD
+from policyengine_core.variables import Variable
 
 from policyengine_uk import Microsimulation, Simulation
+from policyengine_uk.entities import Person
 from policyengine_uk.scenarios.pip_reform import modify_simulation as phase_in_pip
 
 pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
@@ -262,3 +265,77 @@ def test_shipped_pip_callback_retains_its_explicit_periods():
         )
     frame = simulation.to_input_dataframe()
     assert all(f"pip__{year}" in frame for year in years)
+
+
+def test_internal_nested_branch_keeps_unrelated_tax_counterfactual_cache(monkeypatch):
+    situation = family()
+    situation["people"]["parent"]["private_pension_income"] = {YEAR: 20_000}
+    counterfactual = "pension_credit_income_tax_without_employment_income"
+    expected = values(Simulation(situation=situation), counterfactual)
+    assert expected[0] > 0
+    simulation = Simulation(situation=situation)
+    calls = []
+    observed = {}
+    variable = simulation.tax_benefit_system.variables[counterfactual]
+    namespace = variable.get_formula(YEAR).__globals__
+    original = namespace["income_tax_with_inputs_zeroed"]
+
+    def counted(owner, period, variables):
+        calls.append(owner.branch_name)
+        return original(owner, period, variables)
+
+    monkeypatch.setitem(namespace, "income_tax_with_inputs_zeroed", counted)
+
+    class nested_tax_counterfactual_probe(Variable):
+        value_type = float
+        entity = Person
+        definition_period = YEAR_PERIOD
+        label = "Real income tax read through an internal nested branch"
+
+        def formula(person, period, parameters):
+            # These are genuine policy calculations, not fixed test outputs.
+            person(counterfactual, period)
+            owner = person.simulation
+            middle = owner.get_branch("internal_middle")
+            leaf = middle.get_branch("internal_leaf")
+            observed["leaf"] = leaf
+            try:
+                # UC housing does not enter tax on earnings. The inherited
+                # tax counterfactual remains valid through this internal edit.
+                leaf.set_input("uc_housing_costs_element", period, [0])
+                return leaf.calculate(counterfactual, period)
+            finally:
+                del middle.branches["internal_leaf"]
+                del owner.branches["internal_middle"]
+
+    simulation.tax_benefit_system.add_variable(nested_tax_counterfactual_probe)
+    np.testing.assert_allclose(
+        values(simulation, "nested_tax_counterfactual_probe"), expected, atol=0.01
+    )
+    assert observed["leaf"].parent_branch.parent_branch is simulation
+    assert calls == ["default"]
+    assert simulation.branches == {}
+
+
+def test_idle_nested_user_branch_edit_matches_cold_parameter_benchmark():
+    simulation = Simulation(situation=family())
+    parent_net = values(simulation, "household_net_income")
+    middle = simulation.get_branch("user_middle")
+    leaf = middle.get_branch("user_leaf")
+    warm_net = values(leaf, "household_net_income")
+    add_benefit(leaf)
+    expected = benchmark()
+    expected_net = values(expected, "household_net_income")
+    assert expected_net[0] > warm_net[0]
+    for reset in (False, True):
+        if reset:
+            leaf.reset_calculations()
+        for variable in (BENEFIT, "household_net_income"):
+            np.testing.assert_allclose(
+                values(leaf, variable), values(expected, variable), atol=0.01
+            )
+    np.testing.assert_array_equal(
+        values(simulation, "household_net_income"), parent_net
+    )
+    assert simulation.branches["user_middle"] is middle
+    assert middle.branches["user_leaf"] is leaf
diff --git a/policyengine_uk/tests/test_property_income_properties.py b/policyengine_uk/tests/test_property_income_properties.py
index b804dd3172..fc1dc77508 100644
--- a/policyengine_uk/tests/test_property_income_properties.py
+++ b/policyengine_uk/tests/test_property_income_properties.py
@@ -102,6 +102,9 @@ def set_inputs(sim, drawn, uses_allowance=None):
     sim.set_input("property_rental_income", YEAR, receipts)
     sim.set_input("property_finance_costs", YEAR, costs)
     sim.set_input("property_finance_costs_brought_forward", YEAR, brought_forward)
+    if uses_allowance is None:
+        # An automatic draw must not inherit the preceding forced test route.
+        sim.delete_arrays("uses_property_allowance", YEAR)
     sim.reset_calculations()
     if uses_allowance is not None:
         sim.set_input("uses_property_allowance", YEAR, np.full(PEOPLE, uses_allowance))
@@ -335,6 +338,8 @@ def set_couple_inputs(sim, drawn, uses_allowance=None):
     sim.set_input("property_income", YEAR, profit)
     sim.set_input("property_rental_income", YEAR, receipts)
     sim.set_input("property_finance_costs", YEAR, costs)
+    if uses_allowance is None:
+        sim.delete_arrays("uses_property_allowance", YEAR)
     sim.reset_calculations()
     if uses_allowance is not None:
         sim.set_input("uses_property_allowance", YEAR, uses_allowance)
@@ -496,6 +501,8 @@ def set_pensioner_inputs(sim, drawn, uses_allowance=None):
     sim.set_input("property_rental_income", PENSIONER_YEAR, profit + expenses)
     sim.set_input("property_finance_costs", PENSIONER_YEAR, costs)
     sim.set_input("would_claim_pc", PENSIONER_YEAR, np.array(claims_pc))
+    if uses_allowance is None:
+        sim.delete_arrays("uses_property_allowance", PENSIONER_YEAR)
     sim.reset_calculations()
     if uses_allowance is not None:
         sim.set_input("uses_property_allowance", PENSIONER_YEAR, uses_allowance)

From 377af043c2096ddacbe76e1511c8cf421ec048c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 18:00:51 +0000
Subject: [PATCH 5/6] Address modifier input lifecycle review findings

---
 ...fix-scenario-modifier-inputs-2666.fixed.md |   4 +-
 docs/book/usage/scenarios.md                  |  28 ++-
 policyengine_uk/simulation.py                 |   1 +
 .../tests/test_scenario_modifier_inputs.py    | 172 ++++++++++++++++++
 .../marginal_tax_rate_on_capital_gains.py     |  42 +++--
 pyproject.toml                                |   2 +-
 uv.lock                                       |   2 +-
 7 files changed, 221 insertions(+), 30 deletions(-)

diff --git a/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
index d0cf179632..6d54a3afa0 100644
--- a/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
+++ b/changelog.d/fix-scenario-modifier-inputs-2666.fixed.md
@@ -1 +1,3 @@
-Keep inputs stored by scenario modifiers and direct callbacks through resets and all marginal-rate calculations, including input-only variables absent at construction. Refresh dependent values after direct input edits so labour supply responses use the completed inputs, while preserving calculation-time framework writes and independent baselines.
+Retain every surviving explicit input through `reset_calculations`, including ordinary `set_input` calls made after construction outside any modifier or scenario. A reset alone no longer undoes those overrides: use `delete_arrays(variable, period)`, then reset or recalculate to restore the variable's formula or default. Retain these inputs through dynamics, marginal-rate calculations and exports, including input-only variables absent at construction.
+
+Refresh dependent values before the next outermost read after every top-level `set_input` on a built simulation, even without a modifier. Preserve selective invalidation for internal formula and framework writes and keep independent baselines. Use collision-safe temporary branches with cleanup on failure for all three marginal rates, and avoid duplicate cleanup after a successful construction-time modifier. Require PolicyEngine Core 3.33.0 or later for the documented public `drop_computed_arrays` method.
diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md
index 44fbbc32d4..a8665e07d8 100644
--- a/docs/book/usage/scenarios.md
+++ b/docs/book/usage/scenarios.md
@@ -177,19 +177,27 @@ through the earnings step. For a policy that should vary with earnings, use
 parameter changes or a formula reform; setting its calculated value once cannot
 represent that dependence.
 
-The same input retention applies when a callback calls `set_input` directly on
-an already-built simulation, including through `Scenario.apply`. Resets and all
-three marginal-rate calculations retain surviving inputs visible to that
-simulation, including inputs absent at construction. A top-level `set_input`
-also refreshes dependent calculations before the next read, so a callback can
-read Child Benefit, add an amount and then obtain fresh income and response
-values. Inputs stored only on another clone or branch do not become the
-parent's inputs. Formula and framework writes during a calculation keep their
-existing calculation lifecycle.
+This input retention also applies to ordinary `set_input` calls on an
+already-built simulation, including calls outside any modifier or scenario.
+`reset_calculations` retains every surviving explicit input visible to that
+simulation, including inputs absent at construction, and all three marginal-rate
+calculations retain them too. To undo an override, call
+`simulation.delete_arrays(variable, period)`, then reset or recalculate so the
+variable's formula or default runs again. A reset alone no longer undoes an
+input supplied after construction.
+
+Every top-level `set_input` on a built simulation queues dependent-cache cleanup
+before its next outermost read. This includes simulations with no modifier: a
+caller can read Child Benefit, add an amount and then obtain fresh income and
+response values. Inputs stored only on another clone or branch do not become
+the parent's inputs. Formula and framework writes during a calculation keep
+their existing selective invalidation.
 
 Direct edits of the parameter tree require a separate array-cache boundary.
 After calling a parameter's `update` method on a built simulation, call
-`simulation.drop_computed_arrays()` before reading its results. In a wrapper
+`simulation.drop_computed_arrays()` before reading its results. This public
+method requires PolicyEngine Core 3.33.0 or later, the model's minimum version.
+In a wrapper
 that runs a policy callback, late dynamic parameter updates and a dynamic
 callback, clear computed arrays after each operation in that order. This
 preserves the input values and lets the next operation see the updated policy.
diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
index 735ae8032e..e1bab79321 100644
--- a/policyengine_uk/simulation.py
+++ b/policyengine_uk/simulation.py
@@ -442,6 +442,7 @@ def __init__(
                 with self._preserve_stored_inputs():
                     after_data_load.simulation_modifier(self)
                 drop_calculated_values(self)
+                self._external_input_change_pending = False
         self._refresh_capital_gains_projection()
         self._capture_cgt_supplied_inputs()
         self._cgt_warning_enabled = True
diff --git a/policyengine_uk/tests/test_scenario_modifier_inputs.py b/policyengine_uk/tests/test_scenario_modifier_inputs.py
index 9fd686529a..88406fccc7 100644
--- a/policyengine_uk/tests/test_scenario_modifier_inputs.py
+++ b/policyengine_uk/tests/test_scenario_modifier_inputs.py
@@ -416,3 +416,175 @@ def fail_pay_rise(branch, name, *args, **kwargs):
     with pytest.raises(RuntimeError, match="synthetic pay-rise failure"):
         values(simulation, variable)
     assert not simulation.branches
+
+
+def capital_gains_family():
+    situation = family()
+    situation["people"]["parent"]["capital_gains"] = {YEAR: 10_000}
+    situation["people"]["partner"] = {
+        "age": {YEAR: 35},
+        "gender": {YEAR: "MALE"},
+        "employment_income": {YEAR: 20_000},
+        "capital_gains": {YEAR: 20_000},
+    }
+    situation["benunits"]["benunit"]["members"].append("partner")
+    situation["households"]["household"]["members"].append("partner")
+    return situation
+
+
+@pytest.mark.parametrize("branch_name", ["adult_1_cg_rise", "adult_2_cg_rise"])
+def test_capital_gains_rate_preserves_preoccupied_user_branch(branch_name):
+    situation = capital_gains_family()
+    expected = values(
+        Simulation(situation=situation), "marginal_tax_rate_on_capital_gains"
+    )
+    assert (expected[[0, 2]] > 0).all()
+    simulation = Simulation(situation=situation)
+    original_gains = values(simulation, "capital_gains")
+    original_benefit = values(simulation, CHILD_BENEFIT)
+    original_names = list(simulation.input_variables)
+    existing = simulation.get_branch(branch_name)
+    existing.set_input(CHILD_BENEFIT, YEAR, [5_000])
+
+    np.testing.assert_allclose(
+        values(simulation, "marginal_tax_rate_on_capital_gains"),
+        expected,
+        atol=1e-5,
+    )
+
+    assert set(simulation.branches) == {branch_name}
+    assert simulation.branches[branch_name] is existing
+    np.testing.assert_array_equal(values(existing, CHILD_BENEFIT), [5_000])
+    np.testing.assert_array_equal(values(simulation, "capital_gains"), original_gains)
+    np.testing.assert_array_equal(values(simulation, CHILD_BENEFIT), original_benefit)
+    assert simulation.input_variables == original_names
+
+
+@pytest.mark.parametrize(
+    "preoccupied", [False, True], ids=["unoccupied", "preoccupied"]
+)
+def test_capital_gains_rate_failure_preserves_user_branch_and_detaches_owned_scratch(
+    monkeypatch, preoccupied
+):
+    situation = capital_gains_family()
+    expected = values(
+        Simulation(situation=situation), "marginal_tax_rate_on_capital_gains"
+    )
+    assert (expected[[0, 2]] > 0).all()
+    simulation = Simulation(situation=situation)
+    original_gains = values(simulation, "capital_gains")
+    existing = None
+    if preoccupied:
+        existing = simulation.get_branch("adult_1_cg_rise")
+        existing.set_input(CHILD_BENEFIT, YEAR, [5_000])
+    original_branches = dict(simulation.branches)
+    captured = {}
+    set_input = Simulation.set_input
+
+    def fail_gains_rise(branch, name, *args, **kwargs):
+        if branch.parent_branch is simulation and name == "capital_gains":
+            captured["branch"] = branch
+            raise RuntimeError("synthetic capital-gains step failure")
+        return set_input(branch, name, *args, **kwargs)
+
+    with monkeypatch.context() as context:
+        context.setattr(Simulation, "set_input", fail_gains_rise)
+        with pytest.raises(RuntimeError, match="synthetic capital-gains step failure"):
+            values(simulation, "marginal_tax_rate_on_capital_gains")
+
+    assert captured["branch"] is not existing
+    assert simulation.branches == original_branches
+    assert all(
+        branch is not captured["branch"] for branch in simulation.branches.values()
+    )
+    if existing is not None:
+        assert simulation.branches["adult_1_cg_rise"] is existing
+        np.testing.assert_array_equal(values(existing, CHILD_BENEFIT), [5_000])
+    np.testing.assert_array_equal(values(simulation, "capital_gains"), original_gains)
+
+
+def test_after_load_modifier_clears_calculated_values_once(monkeypatch):
+    import policyengine_uk.simulation as uk_simulation
+
+    expected = parameter_benchmark()
+    expected_net = values(expected, "household_net_income")
+    captured = {}
+    cleanup_calls = []
+    drop_calculated_values = uk_simulation.drop_calculated_values
+
+    def counted_cleanup(owner):
+        if owner is captured.get("owner") and captured.get("watch_cleanup"):
+            cleanup_calls.append(owner)
+        return drop_calculated_values(owner)
+
+    def modifier(simulation):
+        captured["owner"] = simulation
+        captured["warm_net"] = values(simulation, "household_net_income")
+        add_benefit(simulation)
+        captured["watch_cleanup"] = True
+
+    monkeypatch.setattr(uk_simulation, "drop_calculated_values", counted_cleanup)
+    simulation = Simulation(
+        situation=family(), scenario=Scenario(simulation_modifier=modifier)
+    )
+    assert len(cleanup_calls) == 1
+    assert (expected_net > captured["warm_net"]).all()
+    for _ in range(2):
+        np.testing.assert_allclose(
+            values(simulation, "household_net_income"), expected_net, atol=0.01
+        )
+        assert len(cleanup_calls) == 1
+
+
+@pytest.mark.parametrize("failure", ["modifier", "cleanup"])
+def test_after_load_cleanup_failure_keeps_pending_write(monkeypatch, failure):
+    import policyengine_uk.simulation as uk_simulation
+
+    expected = parameter_benchmark()
+    expected_benefit = values(expected, CHILD_BENEFIT)
+    expected_net = values(expected, "household_net_income")
+    captured = {}
+    cleanup_calls = []
+    drop_calculated_values = uk_simulation.drop_calculated_values
+
+    def fail_cleanup_once(owner):
+        if owner is captured.get("owner") and captured.get("watch_cleanup"):
+            cleanup_calls.append(owner)
+            if failure == "cleanup" and len(cleanup_calls) == 1:
+                raise RuntimeError("synthetic after-load failure")
+        return drop_calculated_values(owner)
+
+    def modifier(simulation):
+        captured["owner"] = simulation
+        captured["warm_net"] = values(simulation, "household_net_income")
+        add_benefit(simulation)
+        captured["stored_benefit"] = (
+            simulation.get_holder(CHILD_BENEFIT)
+            .get_array(period(YEAR), simulation.branch_name)
+            .copy()
+        )
+        np.testing.assert_allclose(
+            captured["stored_benefit"], expected_benefit, atol=0.01
+        )
+        captured["watch_cleanup"] = True
+        if failure == "modifier":
+            raise RuntimeError("synthetic after-load failure")
+
+    monkeypatch.setattr(uk_simulation, "drop_calculated_values", fail_cleanup_once)
+    with pytest.raises(RuntimeError, match="synthetic after-load failure"):
+        Simulation(situation=family(), scenario=Scenario(simulation_modifier=modifier))
+
+    simulation = captured["owner"]
+    assert simulation._external_input_change_pending
+    np.testing.assert_array_equal(
+        simulation.get_holder(CHILD_BENEFIT).get_array(
+            period(YEAR), simulation.branch_name
+        ),
+        captured["stored_benefit"],
+    )
+    assert (expected_net > captured["warm_net"]).all()
+    np.testing.assert_allclose(
+        values(simulation, "household_net_income"), expected_net, atol=0.01
+    )
+    assert not simulation._external_input_change_pending
+    assert len(cleanup_calls) == (1 if failure == "modifier" else 2)
diff --git a/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py b/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
index 7b6948b849..e9a4f27a11 100644
--- a/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
+++ b/policyengine_uk/variables/gov/hmrc/capital_gains_tax/marginal_tax_rate_on_capital_gains.py
@@ -45,21 +45,29 @@ def formula(person, period, parameters):
         # fraction of what household net income does.
         tax = person("capital_gains_tax", period)
         for adult_index in [1, 2]:
-            alt_simulation = simulation.get_branch(f"adult_{adult_index}_cg_rise")
-            mask = adult_index_values == adult_index
-            retained = retained_input_variables(alt_simulation)
-            for variable in simulation.tax_benefit_system.variables:
-                variable_data = simulation.tax_benefit_system.variables[variable]
-                if variable not in retained and not variable_data.is_input_variable():
-                    alt_simulation.delete_arrays(variable)
-            hold_behavioural_responses(alt_simulation, period, held_responses)
-            alt_simulation.set_input("capital_gains", period, gains + mask * step)
-            alt_person = alt_simulation.person
-            # Float32 rounds the higher gains, so the rise actually stored can
-            # differ from the step by up to half the gains' float32 spacing.
-            gains_rise = alt_person("capital_gains", period) - gains
-            tax_rise = alt_person("capital_gains_tax", period) - tax
-            mtr_values += where(mask, tax_rise / where(mask, gains_rise, 1), 0)
-
-            del simulation.branches[f"adult_{adult_index}_cg_rise"]
+            name = f"adult_{adult_index}_cg_rise"
+            while name in simulation.branches or name == simulation.branch_name:
+                name += "_"
+            alt_simulation = simulation.get_branch(name)
+            try:
+                mask = adult_index_values == adult_index
+                retained = retained_input_variables(alt_simulation)
+                for variable in simulation.tax_benefit_system.variables:
+                    variable_data = simulation.tax_benefit_system.variables[variable]
+                    if (
+                        variable not in retained
+                        and not variable_data.is_input_variable()
+                    ):
+                        alt_simulation.delete_arrays(variable)
+                hold_behavioural_responses(alt_simulation, period, held_responses)
+                alt_simulation.set_input("capital_gains", period, gains + mask * step)
+                alt_person = alt_simulation.person
+                # Float32 rounds the higher gains, so the rise actually stored can
+                # differ from the step by up to half the gains' float32 spacing.
+                gains_rise = alt_person("capital_gains", period) - gains
+                tax_rise = alt_person("capital_gains_tax", period) - tax
+                mtr_values += where(mask, tax_rise / where(mask, gains_rise, 1), 0)
+            finally:
+                if simulation.branches.get(name) is alt_simulation:
+                    del simulation.branches[name]
         return mtr_values
diff --git a/pyproject.toml b/pyproject.toml
index 597f8e5166..d23a63c489 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,7 @@ classifiers = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "policyengine-core>=3.32.27",
+    "policyengine-core>=3.33.0",
     "microdf-python>=1.2.1",
     "pydantic>=2.11.7",
     "tables>=3.10.2",
diff --git a/uv.lock b/uv.lock
index 332a126c50..1651ed0665 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2237,7 +2237,7 @@ requires-dist = [
     { name = "hypothesis", marker = "extra == 'dev'" },
     { name = "jupyter-book", marker = "extra == 'dev'", specifier = ">=2.0.0a0" },
     { name = "microdf-python", specifier = ">=1.2.1" },
-    { name = "policyengine-core", specifier = ">=3.32.27" },
+    { name = "policyengine-core", specifier = ">=3.33.0" },
     { name = "pydantic", specifier = ">=2.11.7" },
     { name = "pytest-cov", marker = "extra == 'dev'" },
     { name = "pytest-xdist", marker = "extra == 'dev'" },

From bb05334596c07d206113caab17d9bdc68b64726d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Sun, 25 Oct 2026 19:33:20 +0000
Subject: [PATCH 6/6] Drain framework invalidation before cloning the baseline

---
 policyengine_uk/simulation.py                 |   4 +
 .../tests/test_scenario_modifier_inputs.py    | 144 ++++++++++++++++++
 2 files changed, 148 insertions(+)

diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
index e1bab79321..be2af7ba1b 100644
--- a/policyengine_uk/simulation.py
+++ b/policyengine_uk/simulation.py
@@ -427,6 +427,10 @@ def __init__(
                 capital_gains_projection=capital_gains_projection,
             )
         else:
+            # Finish framework input invalidation before copying the baseline.
+            if self._external_input_change_pending:
+                drop_calculated_values(self)
+                self._external_input_change_pending = False
             self.baseline = self.clone()
         self.baseline._cgt_warning_enabled = False
         self.baseline._defer_uc_health_warning = False
diff --git a/policyengine_uk/tests/test_scenario_modifier_inputs.py b/policyengine_uk/tests/test_scenario_modifier_inputs.py
index 88406fccc7..59b0fad68b 100644
--- a/policyengine_uk/tests/test_scenario_modifier_inputs.py
+++ b/policyengine_uk/tests/test_scenario_modifier_inputs.py
@@ -588,3 +588,147 @@ def modifier(simulation):
     )
     assert not simulation._external_input_change_pending
     assert len(cleanup_calls) == (1 if failure == "modifier" else 2)
+
+
+def test_default_baseline_clones_after_framework_input_drain(monkeypatch):
+    import policyengine_uk.simulation as uk_simulation
+
+    captured = {}
+    cleanup_calls = []
+    move_values = Simulation.move_values
+    clone = Simulation.clone
+    drop_calculated_values = uk_simulation.drop_calculated_values
+
+    def observe_framework_inputs(owner, donor, target):
+        result = move_values(owner, donor, target)
+        if donor == "bus_fare_spending":
+            captured["owner"] = owner
+            assert owner._external_input_change_pending
+            captured["health_element"] = (
+                owner.get_holder("uc_LCWRA_element")
+                .get_array(period(YEAR), owner.branch_name)
+                .copy()
+            )
+        return result
+
+    def counted_cleanup(owner):
+        if owner is captured.get("owner"):
+            cleanup_calls.append(owner)
+        return drop_calculated_values(owner)
+
+    def observe_clone(owner, *args, **kwargs):
+        if owner is captured.get("owner"):
+            assert not owner._external_input_change_pending
+            assert cleanup_calls == [owner]
+            captured["cloned"] = True
+        return clone(owner, *args, **kwargs)
+
+    monkeypatch.setattr(Simulation, "move_values", observe_framework_inputs)
+    monkeypatch.setattr(Simulation, "clone", observe_clone)
+    monkeypatch.setattr(uk_simulation, "drop_calculated_values", counted_cleanup)
+    simulation = Simulation(situation=family())
+
+    assert captured["cloned"]
+    for owner in (simulation, simulation.baseline):
+        assert not owner._external_input_change_pending
+        np.testing.assert_array_equal(
+            owner.get_holder("employment_income_before_lsr").get_array(
+                period(YEAR), owner.branch_name
+            ),
+            [25_000, 0],
+        )
+        np.testing.assert_array_equal(
+            owner.get_holder("uc_LCWRA_element").get_array(
+                period(YEAR), owner.branch_name
+            ),
+            captured["health_element"],
+        )
+
+
+def test_default_constructor_first_reads_do_not_repeat_framework_drain(monkeypatch):
+    import policyengine_uk.simulation as uk_simulation
+
+    reference = Simulation(situation=family())
+    reference.reset_calculations()
+    reference.baseline.reset_calculations()
+    expected = values(reference, "household_net_income")
+    cleanup_calls = []
+    drop_calculated_values = uk_simulation.drop_calculated_values
+
+    def counted_cleanup(owner):
+        cleanup_calls.append(owner)
+        return drop_calculated_values(owner)
+
+    simulation = Simulation(situation=family())
+    monkeypatch.setattr(uk_simulation, "drop_calculated_values", counted_cleanup)
+    for _ in range(2):
+        for owner in (simulation, simulation.baseline):
+            np.testing.assert_allclose(
+                values(owner, "household_net_income"), expected, atol=0.01
+            )
+    assert cleanup_calls == []
+
+
+def test_default_baseline_drain_failure_retains_pending_inputs(monkeypatch):
+    import policyengine_uk.simulation as uk_simulation
+
+    reference = Simulation(situation=family())
+    reference.reset_calculations()
+    expected_tax = values(reference, "income_tax")
+    captured = {}
+    cloned = []
+    move_values = Simulation.move_values
+    clone = Simulation.clone
+    drop_calculated_values = uk_simulation.drop_calculated_values
+
+    def observe_framework_inputs(owner, donor, target):
+        result = move_values(owner, donor, target)
+        if donor == "bus_fare_spending":
+            captured["owner"] = owner
+            captured["health_element"] = (
+                owner.get_holder("uc_LCWRA_element")
+                .get_array(period(YEAR), owner.branch_name)
+                .copy()
+            )
+        return result
+
+    def fail_terminal_cleanup(owner):
+        if owner is captured.get("owner"):
+            assert owner._external_input_change_pending
+            raise RuntimeError("synthetic terminal framework drain failure")
+        return drop_calculated_values(owner)
+
+    def observe_clone(owner, *args, **kwargs):
+        if owner is captured.get("owner"):
+            cloned.append(owner)
+        return clone(owner, *args, **kwargs)
+
+    with monkeypatch.context() as context:
+        context.setattr(Simulation, "move_values", observe_framework_inputs)
+        context.setattr(Simulation, "clone", observe_clone)
+        context.setattr(uk_simulation, "drop_calculated_values", fail_terminal_cleanup)
+        with pytest.raises(
+            RuntimeError, match="synthetic terminal framework drain failure"
+        ):
+            Simulation(situation=family())
+
+    simulation = captured["owner"]
+    assert simulation._external_input_change_pending
+    assert "baseline" not in simulation.__dict__
+    assert cloned == []
+    np.testing.assert_array_equal(
+        simulation.get_holder("employment_income_before_lsr").get_array(
+            period(YEAR), simulation.branch_name
+        ),
+        [25_000, 0],
+    )
+    np.testing.assert_array_equal(
+        simulation.get_holder("uc_LCWRA_element").get_array(
+            period(YEAR), simulation.branch_name
+        ),
+        captured["health_element"],
+    )
+    np.testing.assert_allclose(
+        values(simulation, "income_tax"), expected_tax, atol=0.01
+    )
+    assert not simulation._external_input_change_pending