diff --git a/lib/functions.php b/lib/functions.php
index 75f51f9..ae5356c 100644
--- a/lib/functions.php
+++ b/lib/functions.php
@@ -18,17 +18,7 @@ function task() {
 }
 
 function group($name, $lambda = null) {
-    $thrown = null;
-    builder()->push_group($name);
-    try {
-        if ($lambda instanceof Closure) $lambda();
-    } catch (\Exception $e) {
-        $thrown = $e;
-    }
-    builder()->pop_group();
-    if ($thrown) {
-        throw $e;
-    }
+    builder()->add_group($name, $lambda);
 }
 
 function before($task, $lambda) {
diff --git a/lib/phake/Builder.php b/lib/phake/Builder.php
index d5dd264..1495952 100644
--- a/lib/phake/Builder.php
+++ b/lib/phake/Builder.php
@@ -45,12 +45,22 @@ public function add_task($name, $work, $deps) {
         $this->assign_description($node);
     }
 
-    public function push_group($name) {
+    public function add_group($name, $lambda = null) {
+        $thrown = null;
+
         $this->target_node = $this->target_node->child_with_name($name);
-    }
 
-    public function pop_group() {
+        try {
+            if ($lambda instanceof \Closure) $lambda($this->application, $this->target_node);
+        } catch (\Exception $e) {
+            $thrown = $e;
+        }
+
         $this->target_node = $this->target_node->get_parent();
+
+        if ($thrown) {
+            throw $e;
+        }
     }
 
     public function before($name, $lambda) {
diff --git a/lib/phake/Node.php b/lib/phake/Node.php
index a84bcb5..a9f447f 100644
--- a/lib/phake/Node.php
+++ b/lib/phake/Node.php
@@ -92,9 +92,9 @@ public function invoke(Application $application) {
             return;
         }
 
-        foreach ($this->before  as $t) $t($application);
-        foreach ($this->lambdas as $t) $t($application);
-        foreach ($this->after   as $t) $t($application);
+        foreach ($this->before  as $t) $t($application, $this);
+        foreach ($this->lambdas as $t) $t($application, $this);
+        foreach ($this->after   as $t) $t($application, $this);
 
         $this->has_run = true;
     }
diff --git a/tests/BuilderTest.php b/tests/BuilderTest.php
index 57e378a..ad18242 100644
--- a/tests/BuilderTest.php
+++ b/tests/BuilderTest.php
@@ -81,4 +81,44 @@ public function testGroups()
 );
         $builder->get_application()->invoke('default');
     }
+
+    public function testArguments()
+    {
+        $builder = new Builder();
+
+        $builder->load_runfile($this->getFixture('arguments.php'));
+
+        $this->expectOutputString(<<<EOF
+first
+phake\Application
+phake\Node
+first:test
+phake\Application
+phake\Node
+
+EOF
+        );
+        $builder->get_application()->invoke('default');
+    }
+
+    public function testBuilderGroupException()
+    {
+        $builder = new Builder();
+
+        try {
+            $builder->add_group('exception', function() {
+                throw new \Exception();
+            });
+            $this->fail('No exception thrown');
+        }
+        catch (Exception $e) {
+            // exception caught. Now check our context is still correct
+        }
+
+        $that = $this;
+        $builder->add_group('ok', function($app, $node) use ($that) {
+            $that->assertInstanceOf('phake\Node', $node);
+            $that->assertEquals('ok', $node->get_name());
+        });
+    }
 }
diff --git a/tests/fixtures/arguments.php b/tests/fixtures/arguments.php
new file mode 100644
index 0000000..f6094be
--- /dev/null
+++ b/tests/fixtures/arguments.php
@@ -0,0 +1,15 @@
+<?php
+
+group('first', function($app, $node) {
+    echo $node->get_name() . "\n";
+    echo get_class($app) . "\n";
+    echo get_class($node) . "\n";
+
+    task('test', function($app, $node) {
+        echo $node->get_name() . "\n";
+        echo get_class($app) . "\n";
+        echo get_class($node) . "\n";
+    });
+});
+
+task('default', 'first:test');