diff --git a/cli/src/cmd/shared.rs b/cli/src/cmd/shared.rs
index eb472ab23..cc969efc6 100644
--- a/cli/src/cmd/shared.rs
+++ b/cli/src/cmd/shared.rs
@@ -13,13 +13,10 @@ pub(super) fn get_task<S: AsRef<str>>(replica: &mut Replica, task_arg: S) -> Fal
     let task_arg = task_arg.as_ref();
 
     // first try treating task as a working-set reference
-    match task_arg.parse::<u64>() {
+    match task_arg.parse::<usize>() {
         Ok(i) => {
-            let mut working_set = replica.working_set().unwrap();
-            if i > 0 && i < working_set.len() as u64 {
-                if let Some(task) = working_set[i as usize].take() {
-                    return Ok(task);
-                }
+            if let Some(task) = replica.get_working_set_task(i)? {
+                return Ok(task);
             }
         }
         Err(_) => {}