diff --git a/lib/absinthe/federation/schema.ex b/lib/absinthe/federation/schema.ex
index 8b6ea7a..d53ad58 100644
--- a/lib/absinthe/federation/schema.ex
+++ b/lib/absinthe/federation/schema.ex
@@ -49,14 +49,11 @@ defmodule Absinthe.Federation.Schema do
     |> Absinthe.Pipeline.for_schema(prototype_schema: schema.__absinthe_prototype_schema__())
     |> Absinthe.Pipeline.upto({Absinthe.Phase.Schema.Validation.Result, pass: :final})
     |> Absinthe.Schema.apply_modifiers(schema)
-
-    # TODO: Due to an issue found with rendering the SDL we had to revert this functionality
-    # https://github.com/DivvyPayHQ/absinthe_federation/issues/28
-    # |> Absinthe.Pipeline.without(__MODULE__.Phase.AddFederatedTypes)
-    # |> Absinthe.Pipeline.insert_before(
-    #   Absinthe.Phase.Schema.ApplyDeclaration,
-    #   __MODULE__.Phase.RemoveResolveReferenceFields
-    # )
+    |> Absinthe.Pipeline.without(__MODULE__.Phase.AddFederatedTypes)
+    |> Absinthe.Pipeline.insert_before(
+      Absinthe.Phase.Schema.ApplyDeclaration,
+      __MODULE__.Phase.RemoveResolveReferenceFields
+    )
   end
 
   @spec to_federated_sdl(schema :: Absinthe.Schema.t()) :: String.t()
diff --git a/lib/absinthe/federation/schema/phase/add_federated_directives.ex b/lib/absinthe/federation/schema/phase/add_federated_directives.ex
index e65a35d..5155472 100644
--- a/lib/absinthe/federation/schema/phase/add_federated_directives.ex
+++ b/lib/absinthe/federation/schema/phase/add_federated_directives.ex
@@ -35,13 +35,16 @@ defmodule Absinthe.Federation.Schema.Phase.AddFederatedDirectives do
   defp maybe_add_key_directive(node, %{key_fields: fields}) when is_binary(fields) do
     directive = Directive.build("key", fields: fields)
 
-    add_directive(node, directive)
+    node
+    |> add_directive(directive)
+    |> mark_referenced()
   end
 
   defp maybe_add_key_directive(node, %{key_fields: fields}) when is_list(fields) do
     fields
     |> Enum.map(&Directive.build("key", fields: &1))
     |> Enum.reduce(node, &add_directive(&2, &1))
+    |> mark_referenced()
   end
 
   defp maybe_add_key_directive(node, _meta), do: node
@@ -83,4 +86,11 @@ defmodule Absinthe.Federation.Schema.Phase.AddFederatedDirectives do
   end
 
   defp add_directive(node, _directive), do: node
+
+  defp mark_referenced(%{__private__: private} = node) do
+    new_private = Keyword.put(private, :__absinthe_referenced__, true)
+    %{node | __private__: new_private}
+  end
+
+  defp mark_referenced(node), do: node
 end
diff --git a/test/absinthe/federation/schema/service_field_test.exs b/test/absinthe/federation/schema/service_field_test.exs
index bab2408..f4a0dd9 100644
--- a/test/absinthe/federation/schema/service_field_test.exs
+++ b/test/absinthe/federation/schema/service_field_test.exs
@@ -146,9 +146,6 @@ defmodule Absinthe.Federation.Schema.ServiceFieldTest do
       assert sdl =~ "@key(fields: \"email\")"
     end
 
-    # TODO: Due to an issue found with rendering the SDL we had to revert this functionality
-    # https://github.com/DivvyPayHQ/absinthe_federation/issues/28
-    @tag :skip
     test "returns sdl with federated types/fields removed" do
       query = """
       {
diff --git a/test/absinthe/federation/schema_test.exs b/test/absinthe/federation/schema_test.exs
index 6bfc6d3..a254b90 100644
--- a/test/absinthe/federation/schema_test.exs
+++ b/test/absinthe/federation/schema_test.exs
@@ -31,9 +31,6 @@ defmodule Absinthe.Federation.SchemaTest do
       assert sdl =~ "type User @extends @key(fields: \"id\") {"
     end
 
-    # TODO: Due to an issue found with rendering the SDL we had to revert this functionality
-    # https://github.com/DivvyPayHQ/absinthe_federation/issues/28
-    @tag :skip
     test "does not render federated types" do
       sdl = Absinthe.Federation.to_federated_sdl(TestSchema)