From f71e8642e9f083ca2b1663fbfca8a32066ac72aa Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Wed, 1 Apr 2026 20:51:57 +0000
Subject: [PATCH] Apply task and play environment vars

---
 executor.go      | 64 ++++++++++++++++++++++++++++++++++++++++++++++++
 executor_test.go | 62 ++++++++++++++++++++++++++++++++++++++++++++++
 modules.go       |  7 ++++++
 3 files changed, 133 insertions(+)

diff --git a/executor.go b/executor.go
index 3a7f3d6..279c416 100644
--- a/executor.go
+++ b/executor.go
@@ -34,6 +34,21 @@ type sshExecutorClient interface {
 	Close() error
 }
 
+// environmentSSHClient wraps another SSH client and prefixes commands with
+// shell exports so play/task environment variables reach remote execution.
+type environmentSSHClient struct {
+	sshExecutorClient
+	prefix string
+}
+
+func (c *environmentSSHClient) Run(ctx context.Context, cmd string) (string, string, int, error) {
+	return c.sshExecutorClient.Run(ctx, c.prefix+cmd)
+}
+
+func (c *environmentSSHClient) RunScript(ctx context.Context, script string) (string, string, int, error) {
+	return c.sshExecutorClient.RunScript(ctx, c.prefix+script)
+}
+
 // Executor runs Ansible playbooks.
 //
 // Example:
@@ -1498,6 +1513,55 @@ func (e *Executor) resolveExpr(expr string, host string, task *Task) string {
 	return "{{ " + expr + " }}" // Return as-is if unresolved
 }
 
+// buildEnvironmentPrefix renders merged play/task environment variables as
+// shell export statements.
+func (e *Executor) buildEnvironmentPrefix(host string, task *Task, play *Play) string {
+	env := make(map[string]string)
+
+	if play != nil {
+		for key, value := range play.Environment {
+			env[key] = value
+		}
+	}
+	if task != nil {
+		for key, value := range task.Environment {
+			env[key] = value
+		}
+	}
+
+	if len(env) == 0 {
+		return ""
+	}
+
+	keys := make([]string, 0, len(env))
+	for key := range env {
+		keys = append(keys, key)
+	}
+	slices.Sort(keys)
+
+	parts := make([]string, 0, len(keys))
+	for _, key := range keys {
+		renderedKey := e.templateString(key, host, task)
+		if renderedKey == "" {
+			continue
+		}
+
+		renderedValue := e.templateString(env[key], host, task)
+		parts = append(parts, sprintf("export %s=%s", renderedKey, shellQuote(renderedValue)))
+	}
+
+	if len(parts) == 0 {
+		return ""
+	}
+
+	return join("; ", parts) + "; "
+}
+
+// shellQuote wraps a string in single quotes for shell use.
+func shellQuote(value string) string {
+	return "'" + replaceAll(value, "'", "'\\''") + "'"
+}
+
 // lookupExprValue resolves the first segment of an expression against the
 // executor, task, and inventory scopes.
 func (e *Executor) lookupExprValue(name string, host string, task *Task) (any, bool) {
diff --git a/executor_test.go b/executor_test.go
index 2790649..3afc373 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -231,6 +231,68 @@ func TestExecutor_RunTaskOnHost_Good_DelegateToUsesDelegatedClient(t *testing.T)
 	assert.Equal(t, 1, mock.commandCount())
 }
 
+func TestExecutor_RunTaskOnHost_Good_EnvironmentMergesForCommand(t *testing.T) {
+	e, mock := newTestExecutorWithMock("host1")
+
+	play := &Play{
+		Environment: map[string]string{
+			"APP_ENV":   "play",
+			"PLAY_ONLY": "from-play",
+		},
+	}
+	task := &Task{
+		Name:   "Environment command",
+		Module: "command",
+		Args: map[string]any{
+			"cmd": `echo "$APP_ENV:$PLAY_ONLY:$TASK_ONLY"`,
+		},
+		Environment: map[string]string{
+			"APP_ENV":   "task",
+			"TASK_ONLY": "from-task",
+		},
+		Register: "env_result",
+	}
+
+	mock.expectCommand(`export APP_ENV='task'; export PLAY_ONLY='from-play'; export TASK_ONLY='from-task'; echo`, "task:from-play:from-task\n", "", 0)
+
+	err := e.runTaskOnHost(context.Background(), "host1", []string{"host1"}, task, play)
+	require.NoError(t, err)
+
+	require.NotNil(t, e.results["host1"]["env_result"])
+	assert.Equal(t, "task:from-play:from-task\n", e.results["host1"]["env_result"].Stdout)
+	assert.True(t, mock.hasExecuted(`export APP_ENV='task'; export PLAY_ONLY='from-play'; export TASK_ONLY='from-task'; echo`))
+}
+
+func TestExecutor_RunTaskOnHost_Good_EnvironmentAppliesToShellScript(t *testing.T) {
+	e, mock := newTestExecutorWithMock("host1")
+
+	play := &Play{
+		Environment: map[string]string{
+			"SHELL_ONLY": "from-play",
+		},
+	}
+	task := &Task{
+		Name:   "Environment shell",
+		Module: "shell",
+		Args: map[string]any{
+			"_raw_params": `echo "$SHELL_ONLY"`,
+		},
+		Environment: map[string]string{
+			"SHELL_ONLY": "from-task",
+		},
+		Register: "shell_env_result",
+	}
+
+	mock.expectCommand(`export SHELL_ONLY='from-task'; echo`, "from-task\n", "", 0)
+
+	err := e.runTaskOnHost(context.Background(), "host1", []string{"host1"}, task, play)
+	require.NoError(t, err)
+
+	require.NotNil(t, e.results["host1"]["shell_env_result"])
+	assert.Equal(t, "from-task\n", e.results["host1"]["shell_env_result"].Stdout)
+	assert.True(t, mock.hasExecuted(`export SHELL_ONLY='from-task'; echo`))
+}
+
 func TestExecutor_RunRole_Good_AppliesRoleTagsToTasks(t *testing.T) {
 	dir := t.TempDir()
 	roleTasks := `---
diff --git a/modules.go b/modules.go
index 6eb2a01..4f22a6c 100644
--- a/modules.go
+++ b/modules.go
@@ -34,6 +34,13 @@ func (e *Executor) executeModule(ctx context.Context, host string, client sshExe
 		defer client.SetBecome(oldBecome, oldUser, oldPass)
 	}
 
+	if prefix := e.buildEnvironmentPrefix(host, task, play); prefix != "" {
+		client = &environmentSSHClient{
+			sshExecutorClient: client,
+			prefix:            prefix,
+		}
+	}
+
 	// Template the args
 	args := e.templateArgs(task.Args, host, task)