Repository: Nukesor/pueue (language: rust)

Issue:
Edit task environment variables
### A detailed description of the feature you would like to see added.

Currently, you can edit three aspects of a scheduled task: the command to be run, the path where the task is run and it's label. 

Tasks have another component: the environment variables under which the command is executed. Can we have an option to edit these?

A stretch goal would be the ability to add, set or remove environment variables _across all tasks or a task group_. This lets you do things like set bandwidth limits on a whole group of tasks where the commands in those tasks honour a bandwidth limit set in an environment variable. Powerful stuff!

### Explain your usecase of the requested feature

Loads of UNIX command change their behaviour based on the environment variables that are present when they start. Sometimes, we need to change what environment variables are present for a given task, because sometimes circumstances change or we made a mistake with the env vars when the task was added.

The ability to alter environment variables across all tasks is another powerful option. This may need to be a separate command (or even a separate FR!).

### Alternatives

I've actively edited the command for tasks to add extra environment variables. This is a poor workaround however, as it is easy to accidentally alter the command itself.

### Additional context

_No response_


Discussion / hints:
It's a bit of an edge-case, but I definitely see how this is a good QoL feature.

I'm thinking about the best way to do this.
My first thought would be to set the key-values directly on the commandline, especially if one wants to set multiple values at the same time (e.g. `pueue edit 0 -e KEY=value -e OTHER_KEY=othervalue`).
The behavior will then be a bit inconsistent, but I feel like editing environment variables in an editor will be pretty cumbersome as there're usually a lot of variables.

Maybe we should add a new subcommand for this? This would also make it quite easy to edit variables for groups or multiple commands, as i feel like those flags would be very confusing on the current `edit` command.
Yes, the more I think about this, the more this sounds like a separate subcommand.
If you like, feel free to start hacking on this :) Your contributions have always been top notch and I believe that you'll make the right design decisions ;D
Just wanted to chime in with another use case. 

I currently schedule jobs on a machine with GPUs using `pueue`. Most times, I need multiple GPUs and set the `CUDA_VISIBLE_DEVICES` argument with a list of GPU IDs for a particular run. To get the queueing effect for that run jobs serially on one set of GPUs, I create a group with the default parallel `1` but manually attach the variable with a bash function,
```shell
function pugrun {
  GPUS=${GPUS:-0}

  pueue group add gpu-${GPUS}

  CUDA_VISIBLE_DEVICES=${GPUS} pueue add -g gpu-${GPUS} -- "${@}"

  unset GPUS
}
```

Editing/adding environment variables to groups would be a great addition! Thank you for the amazing project!
@activatedgeek It sounds like you could make good use of the https://github.com/Nukesor/pueue/wiki/Get-started#load-balancing feature. This would allow you to do automatic balancing on your gpus for you :)
Thank you.

My understanding was that this behavior would require me to change existing code to parse out the `PUEUE_WORKER_ID` and transform it to appropriate values of `CUDA_VISIBLE_DEVICES` (e.g. something like a round robin over the list of GPUs). To avoid going back and change internal code plumbing, it was the easiest to just modify at pueue add time.

But of course, nothing too complicated since it can be similarly plugged in by instead of calling the GPU job directly, calling another bash script that does the transform to call to the GPU job.

Another tricky scenario with the example was this series of steps, say for group `cluster1` with parallel `2` (assuming a machine with 2 GPUs with each job using 1 GPU in a round-robin manner):
```
Queue experiment 1 on cluster 1
Queue experiment 2 on cluster 1
Queue experiment 3 on cluster 1
Scheduled experiment 1 (PUEUE_WORKER_ID = 0 assigned to GPU 0)
Scheduled experiment 2 (PUEUE_WORKER_ID = 1 assigned to GPU 1)
Completed experiment 2
Scheduled experiment 3 (PUEUE_WORKER_ID = 2 assigned to GPU 0)
```

Now experiment 1 and experiment 3 both run on GPU 0, and in many of my cases with large models would lead to an out of memory error on GPU 0. Of course, this could be fixed with slightly smarter scheduling at the expense of more plumbing. This is why I took the easier route of limiting each group to schedule on a fix set of GPU IDs.

Let me know if the scenario makes sense, and if I am missing something simpler here. 
> My understanding was that this behavior would require me to change existing code to parse out the PUEUE_WORKER_ID and transform it to appropriate values of CUDA_VISIBLE_DEVICES (e.g. something like a round robin over the list of GPUs). 

Exactly. That's how it was designed to be used, if the Ids weren't directly mappable to the external resources. (Except the round-robin)

> Now experiment 1 and experiment 3 both run on GPU 0, and in many of my cases with large models would lead to an out of memory error on GPU 0. Of course, this could be fixed with slightly smarter scheduling at the expense of more plumbing. This is why I took the easier route of limiting each group to schedule on a fix set of GPU IDs.

That's actually not how it would work. In your example, the last line would look like this:

```
Completed experiment 2
Scheduled experiment 3 (PUEUE_WORKER_ID = 1 assigned to GPU 1)
```

The ID is not the task id, but rather the internal "worker" or rather "slot" id. A queue with 5 parallel tasks will get 5 slots with ids `0-4`. As soon as a task in a queue finishes, its "slot" is freed and the next task then gets that slot.

If it doesn't work this way, this is definitely a bug! That system was specifically designed to handle your usecase. Machines/pools with a set number of limited resources, to which the tasks are then assigned in a greedy manner.


If the current wiki entry doesn't properly convey this behavior, it would be awesome if you could rephrase it! The wiki is free to be edited by anyone :)
Ahaa! That makes much more sense. Thank you for clarifying. 

Let me run a quick test simulating this scenario, and will report back. I'll update the wiki with this example so that the behavior is more explicit.
Sorry for the delay. I have verified this behavior, and this actually solves my specific problem.

I also realized I missed this:

> There'll never be two concurrent tasks with the same $PUEUE_WORKER_ID in the same group.

In that case, the wiki already explains it precisely, and I will skip the wiki edit.

@mjpieters I'm thinking about adding a `pueue task` command to inspect and eventually manipulate tasks.

The idea came up during https://github.com/Nukesor/pueue/issues/540.

What do you think about a `pueue task env $TASK_ID` command to inspect the environment.
And respective `pueue task env unset $TASK_ID $NAME` and `pueue task env set $TASK_ID $NAME $VALUE` commands?

Is the nesting too big or is this reasonable?
Ping @mjpieters 
I love the nesting, actually. Grouping like this can really help with command clarity. E.g. the `gh` command-line tool for GitHub has 3 levels for most commands.

I would almost say grouping is overdue for pueue; you have commands that operate on either a group or a task, which maybe could do with splitting across command groups, with aliases for backwards compatibility / ease of typing:

* `pueue task` *(aliases: `pueue t`)*
    * `add` *(aliases: `pueue add`)*
    * `remove <TASKIDS>...` *(aliases: `pueue task rm`, `pueue remove`)*
    * `switch <TASKID1> <TASKID2>` *(aliases: `pueue switch`)*
    * `stash <TASKIDS>...` *(aliases: `pueue stash`)*
    * `enqueue [TASKIDS]...` *(aliases: `pueue enqueue`)*
    * `start [TASKIDS]...` *(aliases: `pueue start`)*
    * `restart [TASKIDS]...` *(aliases: `pueue task re`, `pueue restart`)*
    * `pause [TASKIDS]...` *(aliases: `pueue pause`)*
    * `kill [TASKIDS]...` *(aliases: `pueue kill`)*
    * `send <TASKID> <INPUT>` *(aliases: `pueue send`)*
    * `edit <TASKID>` *(aliases: `pueue edit`)*
    * `log [TASKIDS]...` *(aliases: `pueue log`)*
    * `follow [TASKID]` *(aliases: `pueue task fo`, `pueue follow`)*
    * `wait [TASKIDS]...` *(aliases: `pueue wait`)*
    * `env <TASKID>`
        * `set <TASKID> <ENVVAR> <VALUE>`
        * `unset <TASKID> <ENVVARS>...`
* `pueue group` *(aliases: `pueue g`)*
    * `add <GROUP>`
    * `remove <GROUPS>...` *(aliases: `pueue group rm`)*
    * `start [GROUPS]...` *(aliases: `pueue start -g/--group`, `pueue start -a/--all`)*
    * `restart [GROUPS]...` *(aliases: `pueue group re`, `pueue restart -g/--failed-in-group`, `pueue restart -a/--all`)*
    * `pause [GROUPS]...` *(aliases: `pueue pause -g/--group`, `pueue pause -a/--all`)*
    * `kill [GROUPS]...` *(aliases: `pueue kill -g/--group`, `pueue kill -a/--all`)*
    * `status <GROUPS>...` *(aliases: `pueue status -g/--group`)*
    * `wait [GROUPS]...` *(aliases: `pueue wait -g/--group`, `pueue wait -a/--all`)*
    * `clean <GROUPS>...` *(aliases: `pueue group cleanup`, `pueue group clear`, `pueue clean -g/--group`)*
    * `reset [GROUPS]...` *(aliases: `pueue reset -g/--group`)
    * `parallel [GROUP] [PARALLEL_TASKS]` *(aliases: `pueue parallel -g/--group`)*

and the following first-tier top-level commands:

* `pueue status` *(aliases: `pueue`)*
* `pueue clean` *(aliases: `pueue cleanup`, `pueue clear`)*
* `pueue shutdown`
* `pueue completions`
* `pueue help`

Note that for many commands that accepted a single `-g/--group` name the above structure allows for multiple groups to be named. Ditto for `pueue task env unset <ENVVARS>...` to unset multiple environment variables (but I rather name them explicitly and not default to 'all' if you don't name envvars).
I really like the `group` proposal. Especially, since I intuitively tried to use Pueue like that a few times already :D 

I'm not sure about the `task` subcommand though. I feel conflicted about having the `task` subcommand space mirror a big part of the top-level subcommand space. It just feels a bit off to have 2 or more different ways of achieving the same thing.

I guess, this would also be confusing, as I would expect different commands to perform different tasks?
I know that there're some tools out there that already work this way, but I personally always found those a bit unintuitive/confusing to use.

But permanently moving all those commands (`add`, `remove`, etc.) to the `task` subcommand  space would make the CLI less convenient to use. I still assume that most users won't work with groups on a regular basis, so the main interface should stay as convenient as possible.

In general, I really like the idea to restructure the CLI, especially since the next version will be a breaking major version anyway.

I guess I need to think a bit about your proposal and make my mind up about it.
Fair enough; it's why I included aliases for the original top-level commands in there (as well as thinking about backwards compatibility). It could perhaps depend on how the aliases are represented when you run `pueue --help`.

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bd2722c3..66e63a0d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -82,6 +82,7 @@ TLDR: The new task state representation is more verbose but significantly cleane
 - Allow `pueue status` to order tasks by `enqueue_at`. [#554](https://github.com/Nukesor/pueue/issues/554)
 - Added Windows service on Windows to allow a true daemon experience. [#344](https://github.com/Nukesor/pueue/issues/344) [#567](https://github.com/Nukesor/pueue/pull/567)
 - Add `queued_count` and `stashed_count` to callback template variables. This allows users to fire callbacks when whole groups are finished. [#578](https://github.com/Nukesor/pueue/issues/578)
+- Add new subcommand to set or unset environment variables for tasks. [#503](https://github.com/Nukesor/pueue/issues/503)
 
 ### Fixed
 
diff --git a/pueue/src/client/cli.rs b/pueue/src/client/cli.rs
index d530352b..5f88942b 100644
--- a/pueue/src/client/cli.rs
+++ b/pueue/src/client/cli.rs
@@ -10,7 +10,7 @@ use pueue_lib::network::message::Signal;
 
 use super::commands::WaitTargetStatus;
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub enum SubCommand {
     #[command(
         about = "Enqueue a task for execution.\n\
@@ -278,6 +278,12 @@ pub enum SubCommand {
         task_ids: Vec<usize>,
     },
 
+    #[command(about = "Use this to add or remove environment variables from tasks.")]
+    Env {
+        #[command(subcommand)]
+        cmd: EnvCommand,
+    },
+
     #[command(about = "Use this to add or remove groups.\n\
         By default, this will simply display all known groups.")]
     Group {
@@ -492,7 +498,31 @@ https://github.com/Nukesor/pueue/issues/350#issue-1359083118"
     },
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
+pub enum EnvCommand {
+    /// Set a variable for a specific task's environment.
+    Set {
+        /// The id of the task for which the variable should be set.
+        task_id: usize,
+
+        /// The name of the environment variable to set.
+        key: String,
+
+        /// The value of the environment variable to set.
+        value: String,
+    },
+
+    /// Remove a specific variable from a task's environment.
+    Unset {
+        /// The id of the task for which the variable should be set.
+        task_id: usize,
+
+        /// The name of the environment variable to set.
+        key: String,
+    },
+}
+
+#[derive(Parser, Debug, Clone)]
 pub enum GroupCommand {
     /// Add a group by name.
     Add {
diff --git a/pueue/src/client/client.rs b/pueue/src/client/client.rs
index febdf2e4..bef14f04 100644
--- a/pueue/src/client/client.rs
+++ b/pueue/src/client/client.rs
@@ -17,6 +17,8 @@ use crate::client::cli::{CliArguments, ColorChoice, GroupCommand, SubCommand};
 use crate::client::commands::*;
 use crate::client::display::*;
 
+use super::cli::EnvCommand;
+
 /// This struct contains the base logic for the client.
 /// The client is responsible for connecting to the daemon, sending instructions
 /// and interpreting their responses.
@@ -383,7 +385,7 @@ impl Client {
     /// This function is pretty large, but it consists mostly of simple conversions
     /// of [SubCommand] variant to a [Message] variant.
     fn get_message_from_opt(&self) -> Result<Message> {
-        Ok(match &self.subcommand {
+        Ok(match self.subcommand.clone() {
             SubCommand::Add {
                 command,
                 working_directory,
@@ -406,7 +408,7 @@ impl Client {
                 let mut command = command.clone();
                 // The user can request to escape any special shell characters in all parameter strings before
                 // we concatenated them to a single string.
-                if *escape {
+                if escape {
                     command = command
                         .iter()
                         .map(|parameter| shell_escape::escape(Cow::from(parameter)).into_owned())
@@ -418,20 +420,20 @@ impl Client {
                     path,
                     // Catch the current environment for later injection into the task's process.
                     envs: HashMap::from_iter(vars()),
-                    start_immediately: *start_immediately,
-                    stashed: *stashed,
-                    group: group_or_default(group),
-                    enqueue_at: *delay_until,
-                    dependencies: dependencies.to_vec(),
-                    priority: priority.to_owned(),
-                    label: label.clone(),
-                    print_task_id: *print_task_id,
+                    start_immediately,
+                    stashed,
+                    group: group_or_default(&group),
+                    enqueue_at: delay_until,
+                    dependencies,
+                    priority,
+                    label,
+                    print_task_id,
                 }
                 .into()
             }
             SubCommand::Remove { task_ids } => {
                 if self.settings.client.show_confirmation_questions {
-                    self.handle_user_confirmation("remove", task_ids)?;
+                    self.handle_user_confirmation("remove", &task_ids)?;
                 }
                 Message::Remove(task_ids.clone())
             }
@@ -441,10 +443,10 @@ impl Client {
                 all,
                 delay_until,
             } => {
-                let selection = selection_from_params(*all, group, task_ids);
+                let selection = selection_from_params(all, &group, &task_ids);
                 StashMessage {
                     tasks: selection,
-                    enqueue_at: *delay_until,
+                    enqueue_at: delay_until,
                 }
                 .into()
             }
@@ -452,8 +454,8 @@ impl Client {
                 task_id_1,
                 task_id_2,
             } => SwitchMessage {
-                task_id_1: *task_id_1,
-                task_id_2: *task_id_2,
+                task_id_1,
+                task_id_2,
             }
             .into(),
             SubCommand::Enqueue {
@@ -462,10 +464,10 @@ impl Client {
                 all,
                 delay_until,
             } => {
-                let selection = selection_from_params(*all, group, task_ids);
+                let selection = selection_from_params(all, &group, &task_ids);
                 EnqueueMessage {
                     tasks: selection,
-                    enqueue_at: *delay_until,
+                    enqueue_at: delay_until,
                 }
             }
             .into(),
@@ -475,7 +477,7 @@ impl Client {
                 all,
                 ..
             } => StartMessage {
-                tasks: selection_from_params(*all, group, task_ids),
+                tasks: selection_from_params(all, &group, &task_ids),
             }
             .into(),
             SubCommand::Pause {
@@ -485,8 +487,8 @@ impl Client {
                 all,
                 ..
             } => PauseMessage {
-                tasks: selection_from_params(*all, group, task_ids),
-                wait: *wait,
+                tasks: selection_from_params(all, &group, &task_ids),
+                wait,
             }
             .into(),
             SubCommand::Kill {
@@ -497,19 +499,32 @@ impl Client {
                 ..
             } => {
                 if self.settings.client.show_confirmation_questions {
-                    self.handle_user_confirmation("kill", task_ids)?;
+                    self.handle_user_confirmation("kill", &task_ids)?;
                 }
                 KillMessage {
-                    tasks: selection_from_params(*all, group, task_ids),
-                    signal: signal.clone(),
+                    tasks: selection_from_params(all, &group, &task_ids),
+                    signal,
                 }
                 .into()
             }
             SubCommand::Send { task_id, input } => SendMessage {
-                task_id: *task_id,
+                task_id,
                 input: input.clone(),
             }
             .into(),
+            SubCommand::Env { cmd } => Message::from(match cmd {
+                EnvCommand::Set {
+                    task_id,
+                    key,
+                    value,
+                } => EnvMessage::Set {
+                    task_id,
+                    key,
+                    value,
+                },
+                EnvCommand::Unset { task_id, key } => EnvMessage::Unset { task_id, key },
+            }),
+
             SubCommand::Group { cmd, .. } => match cmd {
                 Some(GroupCommand::Add { name, parallel }) => GroupMessage::Add {
                     name: name.to_owned(),
@@ -528,8 +543,8 @@ impl Client {
                 all,
                 ..
             } => {
-                let lines = determine_log_line_amount(*full, lines);
-                let selection = selection_from_params(*all, group, task_ids);
+                let lines = determine_log_line_amount(full, &lines);
+                let selection = selection_from_params(all, &group, &task_ids);
 
                 let message = LogRequestMessage {
                     tasks: selection,
@@ -538,17 +553,13 @@ impl Client {
                 };
                 Message::Log(message)
             }
-            SubCommand::Follow { task_id, lines } => StreamRequestMessage {
-                task_id: *task_id,
-                lines: *lines,
-            }
-            .into(),
+            SubCommand::Follow { task_id, lines } => StreamRequestMessage { task_id, lines }.into(),
             SubCommand::Clean {
                 successful_only,
                 group,
             } => CleanMessage {
-                successful_only: *successful_only,
-                group: group.clone(),
+                successful_only,
+                group,
             }
             .into(),
             SubCommand::Reset { force, groups, .. } => {
@@ -570,9 +581,9 @@ impl Client {
                 group,
             } => match parallel_tasks {
                 Some(parallel_tasks) => {
-                    let group = group_or_default(group);
+                    let group = group_or_default(&group);
                     ParallelMessage {
-                        parallel_tasks: *parallel_tasks,
+                        parallel_tasks,
                         group,
                     }
                     .into()
diff --git a/pueue/src/daemon/network/message_handler/env.rs b/pueue/src/daemon/network/message_handler/env.rs
new file mode 100644
index 00000000..31ddac05
--- /dev/null
+++ b/pueue/src/daemon/network/message_handler/env.rs
@@ -0,0 +1,57 @@
+use pueue_lib::{network::message::*, settings::Settings, state::SharedState};
+
+use crate::{
+    daemon::{network::message_handler::ok_or_failure_message, state_helper::save_state},
+    ok_or_save_state_failure,
+};
+
+/// Invoked on `pueue env`.
+/// Manage environment variables for tasks.
+/// - Set environment variables
+/// - Unset environment variables
+pub fn env(settings: &Settings, state: &SharedState, message: EnvMessage) -> Message {
+    let mut state = state.lock().unwrap();
+
+    let message = match message {
+        EnvMessage::Set {
+            task_id,
+            key,
+            value,
+        } => {
+            let Some(task) = state.tasks.get_mut(&task_id) else {
+                return create_failure_message(format!("No task with id {task_id}"));
+            };
+
+            if !(task.is_queued() || task.is_stashed()) {
+                return create_failure_message("You can only edit stashed or queued tasks");
+            }
+
+            task.envs.insert(key, value);
+
+            create_success_message("Environment variable set.")
+        }
+        EnvMessage::Unset { task_id, key } => {
+            let Some(task) = state.tasks.get_mut(&task_id) else {
+                return create_failure_message(format!("No task with id {task_id}"));
+            };
+
+            if !(task.is_queued() || task.is_stashed()) {
+                return create_failure_message("You can only edit stashed or queued tasks");
+            }
+
+            match task.envs.remove(&key) {
+                Some(_) => create_success_message("Environment variable unset."),
+                None => create_failure_message(format!(
+                    "No environment variable with key '{key}' found."
+                )),
+            }
+        }
+    };
+
+    // Save the state if there were any changes.
+    if let Message::Success(_) = message {
+        ok_or_save_state_failure!(save_state(&state, settings));
+    }
+
+    message
+}
diff --git a/pueue/src/daemon/network/message_handler/group.rs b/pueue/src/daemon/network/message_handler/group.rs
index 5bbfa123..e1eac39f 100644
--- a/pueue/src/daemon/network/message_handler/group.rs
+++ b/pueue/src/daemon/network/message_handler/group.rs
@@ -1,15 +1,21 @@
 use std::collections::BTreeMap;
 
-use pueue_lib::network::message::*;
-use pueue_lib::settings::Settings;
-use pueue_lib::state::{SharedState, PUEUE_DEFAULT_GROUP};
-use pueue_lib::{failure_msg, success_msg};
+use pueue_lib::{
+    failure_msg,
+    network::message::*,
+    settings::Settings,
+    state::{SharedState, PUEUE_DEFAULT_GROUP},
+    success_msg,
+};
 
 use crate::daemon::network::message_handler::ok_or_failure_message;
-use crate::daemon::network::response_helper::ensure_group_exists;
-use crate::daemon::process_handler::initiate_shutdown;
-use crate::daemon::state_helper::save_state;
-use crate::ok_or_save_state_failure;
+use crate::{
+    daemon::{
+        network::response_helper::ensure_group_exists, process_handler::initiate_shutdown,
+        state_helper::save_state,
+    },
+    ok_or_save_state_failure,
+};
 
 /// Invoked on `pueue groups`.
 /// Manage groups.
diff --git a/pueue/src/daemon/network/message_handler/mod.rs b/pueue/src/daemon/network/message_handler/mod.rs
index 24d29c67..0bf39eee 100644
--- a/pueue/src/daemon/network/message_handler/mod.rs
+++ b/pueue/src/daemon/network/message_handler/mod.rs
@@ -12,6 +12,7 @@ mod add;
 mod clean;
 mod edit;
 mod enqueue;
+mod env;
 mod group;
 mod kill;
 mod log;
@@ -34,6 +35,7 @@ pub fn handle_message(message: Message, state: &SharedState, settings: &Settings
         Message::Edit(editable_tasks) => edit::edit(settings, state, editable_tasks),
         Message::EditRequest(task_ids) => edit::edit_request(state, task_ids),
         Message::EditRestore(task_ids) => edit::edit_restore(state, task_ids),
+        Message::Env(message) => env::env(settings, state, message),
         Message::Enqueue(message) => enqueue::enqueue(settings, state, message),
         Message::Group(message) => group::group(settings, state, message),
         Message::Kill(message) => kill::kill(settings, state, message),
diff --git a/pueue/src/daemon/process_handler/finish.rs b/pueue/src/daemon/process_handler/finish.rs
index 292a18a7..3c2716a0 100644
--- a/pueue/src/daemon/process_handler/finish.rs
+++ b/pueue/src/daemon/process_handler/finish.rs
@@ -115,6 +115,7 @@ pub fn handle_finished_tasks(settings: &Settings, state: &mut LockedState) {
 
             task.clone()
         };
+        info!("WTF");
         spawn_callback(settings, state, &task);
 
         if let TaskResult::Failed(_) = result {
diff --git a/pueue_lib/src/network/message.rs b/pueue_lib/src/network/message.rs
index 9da08da1..73146380 100644
--- a/pueue_lib/src/network/message.rs
+++ b/pueue_lib/src/network/message.rs
@@ -66,6 +66,8 @@ pub enum Message {
     /// The client sends the edited details to the daemon.
     Edit(Vec<EditableTask>),
 
+    Env(EnvMessage),
+
     Group(GroupMessage),
     GroupResponse(GroupResponseMessage),
 
@@ -273,6 +275,21 @@ impl EditableTask {
     }
 }
 
+#[derive(PartialEq, Eq, Clone, Debug, Deserialize, Serialize)]
+pub enum EnvMessage {
+    Set {
+        task_id: usize,
+        key: String,
+        value: String,
+    },
+    Unset {
+        task_id: usize,
+        key: String,
+    },
+}
+
+impl_into_message!(EnvMessage, Message::Env);
+
 #[derive(PartialEq, Eq, Clone, Debug, Deserialize, Serialize)]
 pub enum GroupMessage {
     Add {


Test patch:
diff --git a/pueue/tests/client/integration/env.rs b/pueue/tests/client/integration/env.rs
new file mode 100644
index 00000000..9d0ae0e7
--- /dev/null
+++ b/pueue/tests/client/integration/env.rs
@@ -0,0 +1,60 @@
+use anyhow::Result;
+
+use crate::client::helper::*;
+
+/// Set an environment variable an make sure it's there afterwards.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn set_environment() -> Result<()> {
+    let daemon = daemon().await?;
+    let shared = &daemon.settings.shared;
+
+    // Add a stashed task so we can edit it.
+    run_client_command(shared, &["add", "--stashed", "echo $TEST_VARIABLE"])?;
+
+    // Set the environment variable
+    run_client_command(shared, &["env", "set", "0", "TEST_VARIABLE", "thisisatest"])?;
+
+    // Now start the command and wait for it to finish
+    run_client_command(shared, &["enqueue", "0"])?;
+    wait_for_task_condition(shared, 0, |task| task.is_done()).await?;
+
+    let state = get_state(shared).await?;
+    println!("{:#?}", state.tasks[&0].envs);
+
+    // Make sure the environment variable has been set.
+    let output = run_client_command(shared, &["follow", "0"])?;
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert_eq!("thisisatest", stdout.trim());
+
+    Ok(())
+}
+
+/// Set an environment variable an make sure it's there afterwards.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn unset_environment() -> Result<()> {
+    let daemon = daemon().await?;
+    let shared = &daemon.settings.shared;
+
+    // Add a stashed task so we can edit it.
+    run_client_command(shared, &["add", "--stashed", "echo $TEST_VARIABLE"])?;
+
+    // Set the environment variable
+    run_client_command(shared, &["env", "set", "0", "TEST_VARIABLE", "thisisatest"])?;
+
+    // Unset the environment variable again.
+    run_client_command(shared, &["env", "unset", "0", "TEST_VARIABLE"])?;
+
+    // Now start the command and wait for it to finish
+    run_client_command(shared, &["enqueue", "0"])?;
+    wait_for_task_condition(shared, 0, |task| task.is_done()).await?;
+
+    let state = get_state(shared).await?;
+    println!("{:#?}", state.tasks[&0].envs);
+
+    // Make sure the environment variable has been set.
+    let output = run_client_command(shared, &["follow", "0"])?;
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert_eq!("", stdout.trim());
+
+    Ok(())
+}
diff --git a/pueue/tests/client/integration/mod.rs b/pueue/tests/client/integration/mod.rs
index eedfb7a8..8dfd57dc 100644
--- a/pueue/tests/client/integration/mod.rs
+++ b/pueue/tests/client/integration/mod.rs
@@ -1,6 +1,7 @@
 mod completions;
 mod configuration;
 mod edit;
+mod env;
 mod follow;
 mod group;
 mod log;