diff --git a/includes/utils/class-amp-validation-utils.php b/includes/utils/class-amp-validation-utils.php
index d5f4497035d..3bbcbb9baeb 100644
--- a/includes/utils/class-amp-validation-utils.php
+++ b/includes/utils/class-amp-validation-utils.php
@@ -371,6 +371,7 @@ public static function get_source( $callback ) {
 		$type   = isset( $matches[1] ) ? $matches[1] : null;
 		$source = isset( $matches[2] ) ? $matches[2] : null;
 		if ( ( 'plugins' === $type ) || ( 'themes' === $type ) ) {
+			$type = substr( $type, 0, -1 );
 			return compact( 'type', 'source' );
 		}
 		return null;
@@ -410,9 +411,9 @@ public static function wrapped_callback( $callback ) {
 			$output = ob_get_clean();
 
 			if ( ! empty( $output ) ) {
-				printf( '<!--before:%s-->', esc_attr( $callback['source'] ) );
+				printf( '<!--before:%s:%s-->', esc_attr( $callback['type'] ), esc_attr( $callback['source'] ) );
 				echo $output; // WPCS: XSS ok.
-				printf( '<!--after:%s-->', esc_attr( $callback['source'] ) );
+				printf( '<!--after:%s:%s-->', esc_attr( $callback['type'] ), esc_attr( $callback['source'] ) );
 			}
 			return $result;
 		};
diff --git a/tests/test-class-amp-validation-utils.php b/tests/test-class-amp-validation-utils.php
index d64b7683b13..cc929ae741c 100644
--- a/tests/test-class-amp-validation-utils.php
+++ b/tests/test-class-amp-validation-utils.php
@@ -353,16 +353,16 @@ public function test_callback_wrappers() {
 		do_action( $action_no_argument );
 		$output = ob_get_clean();
 		$this->assertContains( '<div></div>', $output );
-		$this->assertContains( '<!--before:amp', $output );
-		$this->assertContains( '<!--after:amp', $output );
+		$this->assertContains( '<!--before:plugin:amp', $output );
+		$this->assertContains( '<!--after:plugin:amp', $output );
 
 		ob_start();
 		do_action( $action_one_argument, $notice );
 		$output = ob_get_clean();
 		$this->assertContains( $notice, $output );
 		$this->assertContains( sprintf( '<div class="notice notice-warning"><p>%s</p></div>', $notice ), $output );
-		$this->assertContains( '<!--before:amp', $output );
-		$this->assertContains( '<!--after:amp', $output );
+		$this->assertContains( '<!--before:plugin:amp', $output );
+		$this->assertContains( '<!--after:plugin:amp', $output );
 
 		ob_start();
 		do_action( $action_two_arguments, $notice, get_the_ID() );
@@ -371,8 +371,8 @@ public function test_callback_wrappers() {
 		self::output_message( $notice, get_the_ID() );
 		$expected_output = ob_get_clean();
 		$this->assertContains( $expected_output, $output );
-		$this->assertContains( '<!--before:amp', $output );
-		$this->assertContains( '<!--after:amp', $output );
+		$this->assertContains( '<!--before:plugin:amp', $output );
+		$this->assertContains( '<!--after:plugin:amp', $output );
 
 		// This action's callback isn't from a plugin, so it shouldn't be wrapped in comments.
 		ob_start();
@@ -397,7 +397,7 @@ public function test_callback_wrappers() {
 	public function test_get_source() {
 		$plugin = AMP_Validation_Utils::get_source( 'amp_after_setup_theme' );
 		$this->assertContains( 'amp', $plugin['source'] );
-		$this->assertEquals( 'plugins', $plugin['type'] );
+		$this->assertEquals( 'plugin', $plugin['type'] );
 		$the_content = AMP_Validation_Utils::get_source( 'the_content' );
 		$this->assertEquals( null, $the_content );
 		$core_function = AMP_Validation_Utils::get_source( 'the_content' );
@@ -426,8 +426,8 @@ public function test_wrapped_callback() {
 
 		$this->assertEquals( 'Closure', get_class( $wrapped_callback ) );
 		$this->assertContains( strval( get_the_ID() ), $output );
-		$this->assertContains( '<!--before:amp', $output );
-		$this->assertContains( '<!--after:amp', $output );
+		$this->assertContains( '<!--before:plugin:amp', $output );
+		$this->assertContains( '<!--after:plugin:amp', $output );
 
 		$callback         = array(
 			'function'      => array( $this, 'get_string' ),