diff --git a/Chickensoft.LogicBlocks.DiagramGenerator.Tests/test_cases/GenericLogicBlock.g.puml b/Chickensoft.LogicBlocks.DiagramGenerator.Tests/test_cases/GenericLogicBlock.g.puml
index ca502b5..57fe062 100644
--- a/Chickensoft.LogicBlocks.DiagramGenerator.Tests/test_cases/GenericLogicBlock.g.puml
+++ b/Chickensoft.LogicBlocks.DiagramGenerator.Tests/test_cases/GenericLogicBlock.g.puml
@@ -1,6 +1,10 @@
 @startuml GenericLogicBlock
-state "GenericLogicBlock State" as Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType___GenericLogicBlock_State {
+state "GenericLogicBlock State" as Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_State {
   state "StateOne" as Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateOne
   state "StateTwo" as Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateTwo
 }
+
+Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateOne --> Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateTwo : InputOne
+Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateTwo --> Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_StateOne : InputTwo
+[*] --> Chickensoft_LogicBlocks_DiagramGenerator_Tests_TestCases_MyGenericType_GenericLogicBlock_State
 @enduml
\ No newline at end of file
diff --git a/Chickensoft.LogicBlocks.DiagramGenerator/src/Diagrammer.cs b/Chickensoft.LogicBlocks.DiagramGenerator/src/Diagrammer.cs
index 2bf50f5..8761f81 100644
--- a/Chickensoft.LogicBlocks.DiagramGenerator/src/Diagrammer.cs
+++ b/Chickensoft.LogicBlocks.DiagramGenerator/src/Diagrammer.cs
@@ -232,7 +232,9 @@ CancellationToken token
 
     // Base state becomes the root
     var root = new LogicBlockGraph(
-      id: CodeService.GetNameFullyQualified(concreteState, concreteState.Name),
+      id: CodeService.GetNameFullyQualifiedWithoutGenerics(
+        concreteState, concreteState.Name
+      ),
       name: concreteState.Name,
       baseId: CodeService.GetNameFullyQualifiedWithoutGenerics(
         concreteState, concreteState.Name
diff --git a/Chickensoft.LogicBlocks.DiagramGenerator/src/utils/SymbolExtensions.cs b/Chickensoft.LogicBlocks.DiagramGenerator/src/utils/SymbolExtensions.cs
index b13aa4b..ec52b21 100644
--- a/Chickensoft.LogicBlocks.DiagramGenerator/src/utils/SymbolExtensions.cs
+++ b/Chickensoft.LogicBlocks.DiagramGenerator/src/utils/SymbolExtensions.cs
@@ -5,10 +5,19 @@ namespace Chickensoft.SourceGeneratorUtils;
 
 public static class SymbolExtensions {
   public static bool InheritsFromOrEquals(
-      this ITypeSymbol type, ITypeSymbol baseType) =>
+      this ITypeSymbol type, INamedTypeSymbol baseType) =>
         type
           .GetBaseTypesAndThis()
-          .Any(t => SymbolEqualityComparer.Default.Equals(t, baseType));
+          .Any(t => SymbolEqualityComparer.Default.Equals(t, baseType)) ||
+          (
+            baseType.IsGenericType &&
+            type.GetBaseTypesAndThis().Any(t =>
+              SymbolEqualityComparer.Default.Equals(
+                t.OriginalDefinition,
+                baseType.OriginalDefinition
+              )
+            )
+          );
 
   private static IEnumerable<ITypeSymbol> GetBaseTypesAndThis(
     this ITypeSymbol? type
diff --git a/Chickensoft.LogicBlocks.DiagramGenerator/src/visitors/ReturnTypeVisitor.cs b/Chickensoft.LogicBlocks.DiagramGenerator/src/visitors/ReturnTypeVisitor.cs
index 0b97337..83a6daf 100644
--- a/Chickensoft.LogicBlocks.DiagramGenerator/src/visitors/ReturnTypeVisitor.cs
+++ b/Chickensoft.LogicBlocks.DiagramGenerator/src/visitors/ReturnTypeVisitor.cs
@@ -18,7 +18,7 @@ public class ReturnTypeVisitor : CSharpSyntaxWalker {
   public INamedTypeSymbol StateType { get; }
   private readonly HashSet<string> _returnTypes = new();
 
-  public ImmutableHashSet<string> ReturnTypes => _returnTypes.ToImmutableHashSet();
+  public ImmutableHashSet<string> ReturnTypes => [.. _returnTypes];
 
   public ReturnTypeVisitor(
     SemanticModel model,