diff --git a/editor/src/messages/tool/common_functionality/graph_modification_utils.rs b/editor/src/messages/tool/common_functionality/graph_modification_utils.rs
index f3404bcd86..ea9ab3bb60 100644
--- a/editor/src/messages/tool/common_functionality/graph_modification_utils.rs
+++ b/editor/src/messages/tool/common_functionality/graph_modification_utils.rs
@@ -396,11 +396,14 @@ impl<'a> NodeGraphLayer<'a> {
 
 	/// Node id of a protonode if it exists in the layer's primary flow
 	pub fn upstream_node_id_from_protonode(&self, protonode_identifier: &'static str) -> Option<NodeId> {
-		self.horizontal_layer_flow().find(move |node_id| {
-			self.network_interface
-				.implementation(node_id, &[])
-				.is_some_and(move |implementation| *implementation == graph_craft::document::DocumentNodeImplementation::proto(protonode_identifier))
-		})
+		self.horizontal_layer_flow()
+			// Take until a different layer is reached
+			.take_while(|&node_id| node_id == self.layer_node || !self.network_interface.is_layer(&node_id, &[]))
+			.find(move |node_id| {
+				self.network_interface
+					.implementation(node_id, &[])
+					.is_some_and(move |implementation| *implementation == graph_craft::document::DocumentNodeImplementation::proto(protonode_identifier))
+			})
 	}
 
 	/// Find all of the inputs of a specific node within the layer's primary flow, up until the next layer is reached.
diff --git a/editor/src/messages/tool/tool_messages/fill_tool.rs b/editor/src/messages/tool/tool_messages/fill_tool.rs
index fe61621373..6bca37c0df 100644
--- a/editor/src/messages/tool/tool_messages/fill_tool.rs
+++ b/editor/src/messages/tool/tool_messages/fill_tool.rs
@@ -175,11 +175,10 @@ mod test_fill {
 		let mut editor = EditorTestUtils::create();
 		editor.new_document().await;
 		editor.drag_tool(ToolType::Rectangle, 0., 0., 100., 100., ModifierKeys::empty()).await;
-		let color = Color::YELLOW;
-		editor.handle_message(ToolMessage::SelectSecondaryColor { color }).await;
+		editor.select_secondary_color(Color::YELLOW).await;
 		editor.click_tool(ToolType::Fill, MouseKeys::LEFT, DVec2::new(2., 2.), ModifierKeys::SHIFT).await;
 		let fills = get_fills(&mut editor).await;
 		assert_eq!(fills.len(), 1);
-		assert_eq!(fills[0].as_solid().unwrap().to_rgba8_srgb(), color.to_rgba8_srgb());
+		assert_eq!(fills[0].as_solid().unwrap().to_rgba8_srgb(), Color::YELLOW.to_rgba8_srgb());
 	}
 }
diff --git a/editor/src/messages/tool/tool_messages/gradient_tool.rs b/editor/src/messages/tool/tool_messages/gradient_tool.rs
index d90d6b2d8d..aed8469650 100644
--- a/editor/src/messages/tool/tool_messages/gradient_tool.rs
+++ b/editor/src/messages/tool/tool_messages/gradient_tool.rs
@@ -516,3 +516,130 @@ impl Fsm for GradientToolFsmState {
 		responses.add(FrontendMessage::UpdateMouseCursor { cursor: MouseCursorIcon::Default });
 	}
 }
+
+#[cfg(test)]
+mod test_gradient {
+	use crate::messages::portfolio::document::{graph_operation::utility_types::TransformIn, utility_types::misc::GroupFolderType};
+	pub use crate::test_utils::test_prelude::*;
+	use glam::DAffine2;
+	use graphene_core::vector::fill;
+	use graphene_std::vector::style::Fill;
+
+	use super::gradient_space_transform;
+
+	async fn get_fills(editor: &mut EditorTestUtils) -> Vec<(Fill, DAffine2)> {
+		let instrumented = editor.eval_graph().await;
+
+		let document = editor.active_document();
+		let layers = document.metadata().all_layers();
+		layers
+			.filter_map(|layer| {
+				let fill = instrumented.grab_input_from_layer::<fill::FillInput<Fill>>(layer, &document.network_interface, &editor.runtime)?;
+				let transform = gradient_space_transform(layer, document);
+				Some((fill, transform))
+			})
+			.collect()
+	}
+
+	#[tokio::test]
+	async fn ignore_artboard() {
+		let mut editor = EditorTestUtils::create();
+		editor.new_document().await;
+		editor.drag_tool(ToolType::Artboard, 0., 0., 100., 100., ModifierKeys::empty()).await;
+		editor.drag_tool(ToolType::Gradient, 2., 2., 4., 4., ModifierKeys::empty()).await;
+		assert!(get_fills(&mut editor).await.is_empty());
+	}
+
+	#[tokio::test]
+	// TODO: remove once https://github.com/GraphiteEditor/Graphite/issues/2444 is fixed
+	#[should_panic]
+	async fn ignore_raster() {
+		let mut editor = EditorTestUtils::create();
+		editor.new_document().await;
+		editor.create_raster_image(Image::new(100, 100, Color::WHITE), Some((0., 0.))).await;
+		editor.drag_tool(ToolType::Gradient, 2., 2., 4., 4., ModifierKeys::empty()).await;
+		assert!(get_fills(&mut editor).await.is_empty());
+	}
+
+	#[tokio::test]
+	async fn simple_draw() {
+		let mut editor = EditorTestUtils::create();
+		editor.new_document().await;
+		editor.drag_tool(ToolType::Rectangle, -5., -3., 100., 100., ModifierKeys::empty()).await;
+		editor.select_primary_color(Color::GREEN).await;
+		editor.select_secondary_color(Color::BLUE).await;
+		editor.drag_tool(ToolType::Gradient, 2., 3., 24., 4., ModifierKeys::empty()).await;
+		let fills = get_fills(&mut editor).await;
+		assert_eq!(fills.len(), 1);
+		let (fill, transform) = fills.first().unwrap();
+		let gradient = fill.as_gradient().unwrap();
+		// Gradient goes from secondary colour to primary colour
+		let stops = gradient.stops.iter().map(|stop| (stop.0, stop.1.to_rgba8_srgb())).collect::<Vec<_>>();
+		assert_eq!(stops, vec![(0., Color::BLUE.to_rgba8_srgb()), (1., Color::GREEN.to_rgba8_srgb())]);
+		assert!(transform.transform_point2(gradient.start).abs_diff_eq(DVec2::new(2., 3.), 1e-10));
+		assert!(transform.transform_point2(gradient.end).abs_diff_eq(DVec2::new(24., 4.), 1e-10));
+	}
+
+	#[tokio::test]
+	async fn snap_simple_draw() {
+		let mut editor = EditorTestUtils::create();
+		editor.new_document().await;
+		editor
+			.handle_message(NavigationMessage::CanvasTiltSet {
+				angle_radians: f64::consts::FRAC_PI_8,
+			})
+			.await;
+		let start = DVec2::new(0., 0.);
+		let end = DVec2::new(24., 4.);
+		editor.drag_tool(ToolType::Rectangle, -5., -3., 100., 100., ModifierKeys::empty()).await;
+		editor.drag_tool(ToolType::Gradient, start.x, start.y, end.x, end.y, ModifierKeys::SHIFT).await;
+		let fills = get_fills(&mut editor).await;
+		let (fill, transform) = fills.first().unwrap();
+		let gradient = fill.as_gradient().unwrap();
+		assert!(transform.transform_point2(gradient.start).abs_diff_eq(start, 1e-10));
+
+		// 15 degrees from horizontal
+		let angle = f64::to_radians(15.);
+		let direction = DVec2::new(angle.cos(), angle.sin());
+		let expected = start + direction * (end - start).length();
+		assert!(transform.transform_point2(gradient.end).abs_diff_eq(expected, 1e-10));
+	}
+
+	#[tokio::test]
+	async fn transformed_draw() {
+		let mut editor = EditorTestUtils::create();
+		editor.new_document().await;
+		editor
+			.handle_message(NavigationMessage::CanvasTiltSet {
+				angle_radians: f64::consts::FRAC_PI_8,
+			})
+			.await;
+		editor.drag_tool(ToolType::Rectangle, -5., -3., 100., 100., ModifierKeys::empty()).await;
+
+		// Group rectangle
+		let group_folder_type = GroupFolderType::Layer;
+		editor.handle_message(DocumentMessage::GroupSelectedLayers { group_folder_type }).await;
+		let metadata = editor.active_document().metadata();
+		let mut layers = metadata.all_layers();
+		let folder = layers.next().unwrap();
+		let rectangle = layers.next().unwrap();
+		assert_eq!(rectangle.parent(metadata), Some(folder));
+		// Transform the group
+		editor
+			.handle_message(GraphOperationMessage::TransformSet {
+				layer: folder,
+				transform: DAffine2::from_scale_angle_translation(DVec2::new(1., 2.), 0., -DVec2::X * 10.),
+				transform_in: TransformIn::Local,
+				skip_rerender: false,
+			})
+			.await;
+
+		editor.drag_tool(ToolType::Gradient, 2., 3., 24., 4., ModifierKeys::empty()).await;
+		let fills = get_fills(&mut editor).await;
+		assert_eq!(fills.len(), 1);
+		let (fill, transform) = fills.first().unwrap();
+		let gradient = fill.as_gradient().unwrap();
+		assert!(transform.transform_point2(gradient.start).abs_diff_eq(DVec2::new(2., 3.), 1e-10));
+		assert!(transform.transform_point2(gradient.end).abs_diff_eq(DVec2::new(24., 4.), 1e-10));
+	}
+}
diff --git a/editor/src/node_graph_executor.rs b/editor/src/node_graph_executor.rs
index 7785a28534..a15b8fa2e4 100644
--- a/editor/src/node_graph_executor.rs
+++ b/editor/src/node_graph_executor.rs
@@ -1,7 +1,10 @@
 use crate::consts::FILE_SAVE_SUFFIX;
 use crate::messages::animation::TimingInformation;
 use crate::messages::frontend::utility_types::{ExportBounds, FileType};
+use crate::messages::portfolio::document::utility_types::document_metadata::LayerNodeIdentifier;
+use crate::messages::portfolio::document::utility_types::network_interface::NodeNetworkInterface;
 use crate::messages::prelude::*;
+use crate::messages::tool::common_functionality::graph_modification_utils::NodeGraphLayer;
 use glam::{DAffine2, DVec2, UVec2};
 use graph_craft::concrete;
 use graph_craft::document::value::{RenderOutput, TaggedValue};
@@ -920,4 +923,13 @@ impl Instrumented {
 
 		Self::downcast::<Input>(dynamic)
 	}
+
+	pub fn grab_input_from_layer<Input: graphene_std::NodeInputDecleration>(&self, layer: LayerNodeIdentifier, network_interface: &NodeNetworkInterface, runtime: &NodeRuntime) -> Option<Input::Result>
+	where
+		Input::Result: Send + Sync + Clone + 'static,
+	{
+		let node_graph_layer = NodeGraphLayer::new(layer, network_interface);
+		let node = node_graph_layer.upstream_node_id_from_protonode(Input::identifier())?;
+		self.grab_protonode_input::<Input>(&vec![node], runtime)
+	}
 }
diff --git a/editor/src/test_utils.rs b/editor/src/test_utils.rs
index 9827bf7d73..164d62377f 100644
--- a/editor/src/test_utils.rs
+++ b/editor/src/test_utils.rs
@@ -217,6 +217,10 @@ impl EditorTestUtils {
 		self.handle_message(Message::Tool(ToolMessage::SelectPrimaryColor { color })).await;
 	}
 
+	pub async fn select_secondary_color(&mut self, color: Color) {
+		self.handle_message(Message::Tool(ToolMessage::SelectSecondaryColor { color })).await;
+	}
+
 	pub async fn create_raster_image(&mut self, image: graphene_core::raster::Image<Color>, mouse: Option<(f64, f64)>) {
 		self.handle_message(PortfolioMessage::PasteImage {
 			name: None,