fix(command): fix try_flatten of ContextChain - #2987 (#2988)

* fixing try_flatten

* matching try_flatten from vanilla code

* fix cargo

* adding test

crates/pumpkin/src/command/context/command_context.rs

@@ -216,17 +216,17 @@ impl<'a> ContextChain<'a> {
     #[must_use]
     pub fn try_flatten(root: &CommandContext<'a>) -> Option<Self> {
         let mut modifiers = Vec::new();
-        let mut current = root;
-
+        let mut current = Arc::new(root.clone());
         loop {
             if let Some(child) = current.get_child() {
-                modifiers.push(child.clone());
-                current = child;
+                let temp_child = child.clone();
+                modifiers.push(current);
+                current = temp_child;
             } else {
                 return current
                     .command
                     .is_some()
-                    .then(|| Self::new(modifiers, Arc::new(current.clone())));
+                    .then(|| Self::new(modifiers, current));
             }
         }
     }
@@ -547,8 +547,6 @@ impl<'a> CommandContextBuilder<'a> {
 
 #[cfg(test)]
 mod test {
-    use std::sync::Arc;
-
     use crate::command::argument_builder::{ArgumentBuilder, CommandArgumentBuilder};
     use crate::command::context::command_context::{
         CommandContext, CommandContextBuilder, ContextChain, ParsedArgument, Stage,
@@ -558,7 +556,11 @@ mod test {
     use crate::command::errors::command_syntax_error::CommandSyntaxError;
     use crate::command::node::dispatcher::{CommandDispatcher, EmptyResultConsumer};
     use crate::command::node::tree::ROOT_NODE_ID;
-    use crate::command::node::{CommandExecutor, CommandExecutorResult, Redirection};
+    use crate::command::node::{
+        CommandExecutor, CommandExecutorResult, RedirectModifier, Redirection,
+    };
+    use pumpkin_util::math::vector3::Vector3;
+    use std::sync::Arc;
 
     struct TenExecutor;
     impl CommandExecutor for TenExecutor {
@@ -710,4 +712,49 @@ mod test {
 
         assert!(ContextChain::try_flatten(&top_context).is_none());
     }
+
+    struct CustomExecutor;
+    impl CommandExecutor for CustomExecutor {
+        fn execute<'a>(&'a self, context: &'a CommandContext) -> CommandExecutorResult<'a> {
+            Box::pin(async move {
+                let source = &context.source;
+                assert_eq!(source.position, Vector3::new(0f64, 10f64, 0f64));
+                Ok(1)
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn multi_stage_modifier_execution() {
+        let mut dispatcher = CommandDispatcher::new();
+        dispatcher.register(
+            CommandArgumentBuilder::new("foo", "A test command").executes(CustomExecutor),
+        );
+        dispatcher.register(
+            CommandArgumentBuilder::new("bar", "Another test command").redirect_with_modifier(
+                Redirection::Root,
+                RedirectModifier::Custom(Arc::new(|context| {
+                    Box::pin(async move {
+                        let mut new_source = context.source.as_ref().clone();
+                        new_source.position = Vector3::new(0f64, 10f64, 0f64);
+                        Ok(vec![Arc::new(new_source)])
+                    })
+                })),
+            ),
+        );
+        let source = Arc::new(CommandSource::dummy());
+        let result = dispatcher.parse_input("bar foo", &source).await;
+        let top_context = result.context.build("bar foo");
+        let chain = ContextChain::try_flatten(&top_context)
+            .expect("The context should have properly flattened, as it has a command to execute");
+        assert_eq!(chain.get_top_context().source.position, Vector3::default());
+        let chain2 = chain
+            .next_stage()
+            .expect("There should have been the next stage");
+        assert!(chain2.next_stage().is_none());
+        let res = chain
+            .execute_all(&source, dispatcher.consumer.as_ref())
+            .await;
+        assert!(res.is_ok_and(|val| val == 1));
+    }
 }