Change SentryAIActivityListener to non-static class to a singleton pattern

Author: alexsohn1126

src/Sentry.Extensions.AI/Extensions/SentryAIExtensions.cs

@@ -53,7 +53,8 @@ public static ChatOptions AddSentryToolInstrumentation(this ChatOptions options)
     [Experimental(DiagnosticId.ExperimentalFeature)]
     public static IChatClient AddSentry(this IChatClient client, Action<SentryAIOptions>? configure = null)
     {
-        SentryAIActivityListener.Init();
+        // The constructor automatically adds the listener, so we can discard the instance
+        _ = new SentryAiActivityListener();
         return new SentryChatClient(client, configure);
     }
 }

src/Sentry.Extensions.AI/SentryAIActivityListener.cs

@@ -5,39 +5,51 @@ namespace Sentry.Extensions.AI;
 /// <summary>
 /// Listens to FunctionInvokingChatClient's Activity
 /// </summary>
-internal static class SentryAIActivityListener
+internal class SentryAiActivityListener
 {
-    private static volatile IHub Hub = HubAdapter.Instance;
+    private static ActivityListener? Instance;
 
     /// <summary>
-    /// Sentry's <see cref="ActivityListener"/> to tap into function invocation's Activity
+    /// Initializes Sentry's <see cref="ActivityListener"/> to tap into FunctionInvokingChatClient's Activity
     /// </summary>
-    private static readonly ActivityListener FICCListener = new()
+    /// <param name="hub">Optional IHub instance to use. If not provided, HubAdapter.Instance will be used.</param>
+    public SentryAiActivityListener(IHub? hub = null)
     {
-        ShouldListenTo = source => source.Name.StartsWith(SentryAIConstants.SentryActivitySourceName),
-        Sample = (ref ActivityCreationOptions<ActivityContext> options) =>
-            SentryAIConstants.FICCActivityNames.Contains(options.Name) ?
-                ActivitySamplingResult.AllDataAndRecorded : ActivitySamplingResult.None,
-        ActivityStarted = activity =>
+        if (Instance != null)
         {
-            var agentSpan = Hub.StartSpan(SentryAIConstants.SpanAttributes.InvokeAgentOperation, SentryAIConstants.SpanAttributes.InvokeAgentDescription);
-            activity.SetFused(SentryAIConstants.SentryFICCSpanAttributeName, agentSpan);
-        },
-        ActivityStopped = activity =>
-        {
-            var agentSpan = activity.GetFused<ISpan>(SentryAIConstants.SentryFICCSpanAttributeName);
-            // Don't pass in OK status in case there was an exception
-            agentSpan?.Finish();
+            return;
         }
-    };
+
+        var currHub = hub ?? HubAdapter.Instance;
+        Instance = new ActivityListener
+        {
+            ShouldListenTo = source => source.Name.StartsWith(SentryAIConstants.SentryActivitySourceName),
+            Sample = (ref ActivityCreationOptions<ActivityContext> options) =>
+                SentryAIConstants.FICCActivityNames.Contains(options.Name)
+                    ? ActivitySamplingResult.AllDataAndRecorded
+                    : ActivitySamplingResult.None,
+            ActivityStarted = activity =>
+            {
+                var agentSpan = currHub.StartSpan(SentryAIConstants.SpanAttributes.InvokeAgentOperation,
+                    SentryAIConstants.SpanAttributes.InvokeAgentDescription);
+                activity.SetFused(SentryAIConstants.SentryFICCSpanAttributeName, agentSpan);
+            },
+            ActivityStopped = activity =>
+            {
+                var agentSpan = activity.GetFused<ISpan>(SentryAIConstants.SentryFICCSpanAttributeName);
+                // Don't pass in OK status in case there was an exception
+                agentSpan?.Finish();
+            }
+        };
+        ActivitySource.AddActivityListener(Instance);
+    }
 
     /// <summary>
-    /// Initializes Sentry's <see cref="ActivityListener"/> to tap into FunctionInvokingChatClient's Activity
+    /// Dispose the singleton instance (for testing purposes mostly)
     /// </summary>
-    /// <param name="hub">Optional IHub instance to use. If not provided, HubAdapter.Instance will be used.</param>
-    internal static void Init(IHub? hub = null)
+    internal static void Dispose()
     {
-        Hub = hub ?? HubAdapter.Instance;
-        ActivitySource.AddActivityListener(FICCListener);
+        Instance?.Dispose();
+        Instance = null;
     }
 }

test/Sentry.Extensions.AI.Tests/SentryAIActivityListenerTests.cs

@@ -16,11 +16,17 @@ public Fixture()
 
     private readonly Fixture _fixture = new();
 
+    public SentryAIActivityListenerTests()
+    {
+        // Dispose ActivityListener before each test, otherwise the singleton instance will persist between tests
+        SentryAiActivityListener.Dispose();
+    }
+
     [Fact]
     public void Init_AddsActivityListenerToActivitySource()
     {
         // Act
-        SentryAIActivityListener.Init();
+        _ = new SentryAiActivityListener();
 
         // Assert
         Assert.True(SentryAIActivitySource.Instance.HasListeners());
@@ -31,7 +37,7 @@ public void Init_AddsActivityListenerToActivitySource()
     public void ShouldListenTo_ReturnsTrueForSentryActivitySource(string sourceName)
     {
         // Arrange
-        SentryAIActivityListener.Init(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
         var activitySource = new ActivitySource(sourceName);
 
         // Act
@@ -51,7 +57,7 @@ public void ShouldListenTo_ReturnsTrueForSentryActivitySource(string sourceName)
     public void ShouldListenTo_ReturnsFalseForNonSentryActivitySource()
     {
         // Arrange
-        SentryAIActivityListener.Init(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
         var activitySource = new ActivitySource("Other.ActivitySource");
 
         // Act & Assert
@@ -65,12 +71,10 @@ public void ShouldListenTo_ReturnsFalseForNonSentryActivitySource()
 
     [Theory]
     [InlineData("orchestrate_tools")]
-    [InlineData("FunctionInvokingChatClient.GetResponseAsync")]
-    [InlineData("FunctionInvokingChatClient")]
     public void Sample_ReturnsAllDataAndRecordedForFICCActivityNames(string activityName)
     {
         // Arrange
-        SentryAIActivityListener.Init(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
 
         // Act
         using var activity = SentryAIActivitySource.Instance.StartActivity(activityName);
@@ -90,9 +94,9 @@ public void Sample_ReturnsAllDataAndRecordedForFICCActivityNames(string activity
     public void Init_MultipleCalls_NoDuplicateListener_StartsOnlyOneTransaction()
     {
         // Arrange
-        SentryAIActivityListener.Init(_fixture.Hub);
-        SentryAIActivityListener.Init(_fixture.Hub);
-        SentryAIActivityListener.Init(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
+        _ = new SentryAiActivityListener(_fixture.Hub);
 
         // Act
         using var activity = SentryAIActivitySource.Instance.StartActivity(SentryAIConstants.FICCActivityNames[0]);