In a Jenkins Pipeline, is it possible to set env variables to agents that are visible across stages without altering the global env object?

In a Jenkins pipeline, the global env object can be used to set environment variables from most execution contexts, and variables there will be visible among all the agents, possibly creating unexpected race conditions when parallel node execution happens. What I would like to do is configure variables that will be set only in the execution context of the currently running agent, without altering properties returned by the env object in other ones, like the following:

pipeline {
    agent none

    stages {
        stage('Stage1AgentA') {
            agent {
                label "AgentA"
            }
            steps {
               setAgentEnvironmentVariable('MY_VAR', 'MyValue')
               sh 'echo "MY_VAR: $MY_VAR"'      // Should print "MyValue"
               echo "env.MY_VAR: ${env.MY_VAR}" // Should print "MyValue"
            }
        }
        stage('Stage2AgentA') {
            agent {
                label "AgentA"
            }
            steps {
               sh 'echo "MY_VAR: $MY_VAR"'      // Should print "MyValue"
               echo "env.MY_VAR: ${env.MY_VAR}" // Should print "MyValue"
            }
        }
        stage('StageAgentB') {
            agent {
                label "AgentB"
            }
            steps {
               sh 'echo "MY_VAR: $MY_VAR"'      // Should print ""
               echo "env.MY_VAR: ${env.MY_VAR}" // Should print "null"
            }
        }
    }
}

def setAgentEnvironmentVariable(String name, String value)
{
    // Any restricted API is admissible
    // ..
}

I'm already aware that a recommended solution for this use case would be using the withEnv basic step, but I would like to craft a more advanced solution using the restricted Jenkins API.

Using the Jenkins restricted API, it's possible to create an utility method that will add the variable only in the agent execution context. The method can be stored in a Jenkins shared library as groovy code as it follows:

import hudson.EnvVars;
import hudson.model.Run;
import hudson.model.EnvironmentContributingAction;
import org.jenkinsci.plugins.workflow.cps.EnvActionImpl;

/** Set the given environment variable in the executing agent,
 * without altering the global pipeline environment (That is "env.VARNAME").
 * The variables set this way are persistent across stages. Throws when
 * executed in the controller
 */
def setAgentEnvironmentVariable(String name, String value)
{
    // Determine if I'm running in the controller
    def nodeName = env.NODE_NAME;
    if (nodeName == null || nodeName == 'master' || nodeName == 'built-in')
        throw new Exception("The method is not supported when running in a controller built-in node");

    def action = currentBuild.rawBuild.getAction(NodeVariablesAction.class);
    if (action == null)
    {
        action = new NodeVariablesAction();
        currentBuild.rawBuild.addAction(action);
    }

    action.addNodeVariable(nodeName, name, value);
}

class NodeVariablesAction implements EnvironmentContributingAction, Serializable
{
    private static final long serialVersionUID = 1L;

    private HashMap<String, ArrayList<String[]>> nodeVariables;
    private boolean isSetting;

    public NodeVariablesAction()
    {
        nodeVariables = new HashMap<String, ArrayList<String[]>>();
    }

    public void addNodeVariable(String nodeName, String key, String value)
    {
        def variables = nodeVariables.get(nodeName);
        if (variables == null)
        {
            variables = new ArrayList<String[]>();
            nodeVariables.put(nodeName, variables);
        }
        
        def variable = new String[2];
        variable[0] = key;
        variable[1] = value;
        variables.add(variable);
    }

    @NonCPS
    public void buildEnvironment(Run build, EnvVars envVars)
    {
        if (isSetting)
            return;

        // Retrieve the global 'env' object (https://www.jenkins.io/doc/book/pipeline/getting-started/#global-variable-reference)
        // from the build and get the 'NODE_NAME' from there. This is the only working
        // way found, so far. NOTE: Avoid a stack overflow by setting a sentinel that
        // will return from this method in subsequent recursions
        isSetting = true;
        def env = EnvActionImpl.forRun(build);
        def currentNodeName = env.getProperty('NODE_NAME');
        isSetting = false;

        def variables = nodeVariables.get(currentNodeName);
        if (variables == null)
            return;

        for (def variable in variables)
            envVars.put(variable[0], variable[1]);
    }

    @NonCPS
    public String getDisplayName() { return null; }

    @NonCPS
    public String getIconFileName() { return null; }

    @NonCPS
    public String getUrlName() { return null; }
}