diff --git a/lib/src/dartdoc.dart b/lib/src/dartdoc.dart
index a2e4d92ff5..ccbcb28c8f 100644
--- a/lib/src/dartdoc.dart
+++ b/lib/src/dartdoc.dart
@@ -225,7 +225,6 @@ class Dartdoc {
     if (config.showStats) {
       logInfo(runtimeStats.buildReport());
     }
-    await packageBuilder.dispose();
     return DartdocResults(config.topLevelPackageMeta, packageGraph, _outputDir);
   }
 
diff --git a/lib/src/model/inheriting_container.dart b/lib/src/model/inheriting_container.dart
index 65c21dd707..a756d9974a 100644
--- a/lib/src/model/inheriting_container.dart
+++ b/lib/src/model/inheriting_container.dart
@@ -19,9 +19,10 @@ import 'package:meta/meta.dart';
 /// Note that [Constructor]s are not considered to be modifiers so a
 /// [hasModifiers] override is not necessary for this mixin.
 mixin Constructable implements InheritingContainer {
-  late final List<Constructor> constructors = element.constructors
-      .map((e) => getModelFor(e, library) as Constructor)
-      .toList(growable: false);
+  late final List<Constructor> constructors =
+      (element.augmented?.constructors ?? element.constructors)
+          .map((e) => getModelFor(e, library) as Constructor)
+          .toList(growable: false);
 
   @override
   late final List<Constructor> publicConstructorsSorted =
@@ -153,12 +154,12 @@ abstract class InheritingContainer extends Container
 
     // The mapping of all of the inherited element names to their _concrete_
     // implementation element.
-    var concreteInheritanceMap =
-        packageGraph.inheritanceManager.getInheritedConcreteMap2(element);
+    var concreteInheritanceMap = packageGraph.inheritanceManager
+        .getInheritedConcreteMap2(element.augmentedDeclarationOrSelf);
     // The mapping of all inherited element names to the nearest inherited
     // element that they resolve to.
-    var inheritanceMap =
-        packageGraph.inheritanceManager.getInheritedMap2(element);
+    var inheritanceMap = packageGraph.inheritanceManager
+        .getInheritedMap2(element.augmentedDeclarationOrSelf);
 
     var inheritanceChainElements =
         inheritanceChain.map((c) => c.element).toList(growable: false);
@@ -256,9 +257,10 @@ abstract class InheritingContainer extends Container
   }();
 
   @override
-  late final List<Method> declaredMethods = element.methods
-      .map((e) => getModelFor(e, library) as Method)
-      .toList(growable: false);
+  late final List<Method> declaredMethods =
+      (element.augmented?.methods ?? element.methods)
+          .map((e) => getModelFor(e, library) as Method)
+          .toList(growable: false);
 
   @override
   late final List<TypeParameter> typeParameters = element.typeParameters
@@ -594,6 +596,13 @@ mixin MixedInTypes on InheritingContainer {
       mixedInTypes.wherePublic;
 }
 
+extension on InterfaceElement {
+  /// This element's augmented declaration, or, if there is none, then just this
+  /// element itself.
+  InterfaceElement get augmentedDeclarationOrSelf =>
+      augmented?.declaration ?? this;
+}
+
 extension on InterfaceElement {
   bool get isDartCoreObject => name == 'Object' && library.name == 'dart.core';
 }
diff --git a/lib/src/model/package_builder.dart b/lib/src/model/package_builder.dart
index 5629fced1b..085c46e274 100644
--- a/lib/src/model/package_builder.dart
+++ b/lib/src/model/package_builder.dart
@@ -38,8 +38,6 @@ abstract class PackageBuilder {
   // Builds package graph to be used by documentation generator.
   Future<PackageGraph> buildPackageGraph();
 
-  Future<void> dispose();
-
   /// The `include-external` option is deprecated, so we track whether it was
   /// used, to report it.
   bool get includeExternalsWasSpecified;
@@ -119,7 +117,11 @@ class PubPackageBuilder implements PackageBuilder {
 
     logDebug('${DateTime.now()}: Initializing package graph...');
     runtimeStats.startPerfTask('initializePackageGraph');
-    await newGraph.initializePackageGraph();
+    try {
+      await newGraph.initializePackageGraph();
+    } finally {
+      await _dispose();
+    }
     runtimeStats.endPerfTask();
 
     runtimeStats.startPerfTask('initializeCategories');
@@ -129,8 +131,7 @@ class PubPackageBuilder implements PackageBuilder {
     return newGraph;
   }
 
-  @override
-  Future<void> dispose() async {
+  Future<void> _dispose() async {
     // Shutdown macro support.
     await _contextCollection.dispose();
   }