From e47923b025835c9d182a206c491e7f8baa5a55d9 Mon Sep 17 00:00:00 2001
From: bd_ <bd_@nadena.dev>
Date: Thu, 28 Nov 2024 14:56:41 -0800
Subject: [PATCH] feat: add Owner field to IExtensionContext

When set, errors that occur while activating or deactivating the context will
be properly attributed to the owning plugin.

Closes: #463
---
 CHANGELOG.md                         |  2 ++
 Editor/API/BuildContext.cs           | 16 ++++++++++++++++
 Editor/API/IExtensionContext.cs      | 10 +++++++++-
 Editor/ErrorReporting/ErrorReport.cs |  9 +++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0581577..e93754b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [#473] Added `BuildContext.SetEnableUVDistributionRecalculation` to allow opting out from the automatic call to
   `Mesh.RecalculateUVDistributionMetrics` on generated meshes.
 - [#478] Added `ProfilerScope` API
+- [#480] Added `IExtensionContext.Owner` API. Setting this property will allow errors to be correctly attributed to the
+  plugin that contains an extension context. 
 
 ### Fixed
 
diff --git a/Editor/API/BuildContext.cs b/Editor/API/BuildContext.cs
index b77e70f..69bc055 100644
--- a/Editor/API/BuildContext.cs
+++ b/Editor/API/BuildContext.cs
@@ -312,10 +312,18 @@ public void DeactivateExtensionContext(Type t)
                     {
                         var ctx = _activeExtensions[t];
                         Profiler.BeginSample("NDMF Deactivate: " + t);
+
+                        using var scope = ErrorReport.WithContext(ctx);
                         try
                         {
                             ctx.OnDeactivate(this);
                         }
+                        catch (Exception e)
+                        {
+                            // ensure we report the exception while the error report context is set
+                            ErrorReport.ReportException(e);
+                            return;
+                        }
                         finally
                         {
                             Profiler.EndSample();
@@ -405,10 +413,18 @@ public IExtensionContext ActivateExtensionContext(Type ty)
                     if (!_activeExtensions.ContainsKey(ty))
                     {
                         Profiler.BeginSample("NDMF Activate: " + ty);
+
+                        using var scope = ErrorReport.WithContext(ctx);
                         try
                         {
                             ctx.OnActivate(this);
                         }
+                        catch (Exception e)
+                        {
+                            // ensure we report the exception while the error report context is set
+                            ErrorReport.ReportException(e);
+                            return null;
+                        }
                         finally
                         {
                             Profiler.EndSample();
diff --git a/Editor/API/IExtensionContext.cs b/Editor/API/IExtensionContext.cs
index 190c3ac..cbb9395 100644
--- a/Editor/API/IExtensionContext.cs
+++ b/Editor/API/IExtensionContext.cs
@@ -1,4 +1,6 @@
-using System;
+#nullable enable
+
+using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 
@@ -20,6 +22,12 @@ public interface IExtensionContext
         /// </summary>
         /// <param name="context"></param>
         void OnDeactivate(BuildContext context);
+        
+        /// <summary>
+        /// Return the plugin owning this extension context. Implementing this API is optional for backwards
+        /// compatibility, but is encouraged as it will provide better error messaging.
+        /// </summary>
+        PluginBase? Owner => null;
     }
 
     internal static class ExtensionContextUtil
diff --git a/Editor/ErrorReporting/ErrorReport.cs b/Editor/ErrorReporting/ErrorReport.cs
index 7a94179..94863f1 100644
--- a/Editor/ErrorReporting/ErrorReport.cs
+++ b/Editor/ErrorReporting/ErrorReport.cs
@@ -308,6 +308,15 @@ internal IDisposable WithContext(PluginBase thePlugin)
             CurrentContext.Plugin = thePlugin;
             return scope;
         }
+
+        internal IDisposable WithContext(IExtensionContext theExtension)
+        {
+            var scope = new RestoreContextScope(this);
+            
+            if (theExtension.Owner != null) CurrentContext.Plugin = theExtension.Owner;
+
+            return scope;
+        }
         
         internal IDisposable WithContextPassName(string name)
         {