Removed static dependencies

Author: jamescrosswell

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

@@ -51,10 +51,15 @@ public static ChatOptions AddSentryToolInstrumentation(this ChatOptions options)
     /// <param name="configure">The <see cref="SentryAIOptions"/> configuration</param>
     /// <returns>The instrumented <see cref="IChatClient"/></returns>
     [Experimental(DiagnosticId.ExperimentalFeature)]
-    public static IChatClient AddSentry(this IChatClient client, Action<SentryAIOptions>? configure = null)
+    public static IChatClient AddSentry(this IChatClient client, Action<SentryAIOptions>? configure = null) =>
+        AddSentry(client, SentryAiActivityListener.Instance, configure);
+
+    /// <summary>
+    /// Internal overload for testing
+    /// </summary>
+    internal static IChatClient AddSentry(this IChatClient client, ActivityListener listener, Action<SentryAIOptions>? configure = null)
     {
-        // The constructor automatically adds the listener, so we can discard the instance
-        _ = new SentryAiActivityListener();
+        ActivitySource.AddActivityListener(listener);
         return new SentryChatClient(client, configure);
     }
 }

src/Sentry.Extensions.AI/SentryAIActivityListener.cs

@@ -5,6 +5,10 @@ namespace Sentry.Extensions.AI;
 /// <summary>Sentry's <see cref="ActivitySource"/> to be used in <see cref="IChatClient"/></summary>
 internal static class SentryAIActivitySource
 {
-    /// <summary>Sentry's <see cref="ActivitySource"/> to be used in <see cref="IChatClient"/></summary>
-    internal static ActivitySource Instance { get; } = new(SentryAIConstants.SentryActivitySourceName);
+    internal const string SentryActivitySourceName = "Sentry.AgentMonitoring";
+
+    private static readonly Lazy<ActivitySource> LazyInstance = new(CreateSource);
+    internal static ActivitySource Instance => LazyInstance.Value;
+
+    internal static ActivitySource CreateSource() => new(SentryActivitySourceName);
 }

src/Sentry.Extensions.AI/SentryAIActivitySource.cs

@@ -9,11 +9,6 @@ internal static class SentryAIConstants
     internal static readonly string[] FICCActivityNames =
         ["orchestrate_tools", "FunctionInvokingChatClient.GetResponseAsync", "FunctionInvokingChatClient"];
 
-    /// <summary>
-    /// The string we use to identify our <see cref="ActivitySource"/>.
-    /// </summary>
-    internal const string SentryActivitySourceName = "Sentry.AgentMonitoring";
-
     /// <summary>
     /// The string we use to retrieve a Sentry span from the <see cref="Activity"/> using a Fused property
     /// </summary>

src/Sentry.Extensions.AI/SentryAIConstants.cs

@@ -0,0 +1,46 @@
+using Sentry.Internal;
+
+namespace Sentry.Extensions.AI;
+
+/// <summary>
+/// Listens to FunctionInvokingChatClient's Activity
+/// </summary>
+internal static class SentryAiActivityListener
+{
+    /// <summary>
+    /// Singleton used outside of testing
+    /// </summary>
+    private static readonly Lazy<ActivityListener> LazyInstance = new(() => CreateListener());
+    internal static readonly ActivityListener Instance = LazyInstance.Value;
+
+    /// <summary>
+    /// Initializes Sentry's <see cref="ActivityListener"/> to tap into FunctionInvokingChatClient's Activity
+    /// </summary>
+    public static ActivityListener CreateListener(IHub? hub = null)
+    {
+        hub ??= HubAdapter.Instance;
+
+        var listener = new ActivityListener
+        {
+            ShouldListenTo = source => source.Name.StartsWith(SentryAIActivitySource.SentryActivitySourceName),
+            Sample = (ref ActivityCreationOptions<ActivityContext> options) =>
+                SentryAIConstants.FICCActivityNames.Contains(options.Name)
+                    ? ActivitySamplingResult.AllDataAndRecorded
+                    : ActivitySamplingResult.None,
+            ActivityStarted = activity =>
+            {
+                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();
+            }
+        };
+        ActivitySource.AddActivityListener(listener);
+        return listener;
+    }
+}

src/Sentry.Extensions.AI/SentryAiActivityListener.cs

@@ -4,11 +4,20 @@ namespace Sentry.Extensions.AI;
 
 internal sealed class SentryChatClient : DelegatingChatClient
 {
+    private readonly ActivitySource _activitySource;
     private readonly HubAdapter _hub = HubAdapter.Instance;
     private readonly SentryAIOptions _sentryAIOptions;
 
-    public SentryChatClient(IChatClient client, Action<SentryAIOptions>? configure = null) : base(client)
+    public SentryChatClient(IChatClient client, Action<SentryAIOptions>? configure = null) : this(null, client, configure)
     {
+    }
+
+    /// <summary>
+    /// Internal ovverride for testing
+    /// </summary>
+    internal SentryChatClient(ActivitySource? activitySource, IChatClient client, Action<SentryAIOptions>? configure = null) : base(client)
+    {
+        _activitySource = activitySource ?? SentryAIActivitySource.Instance;
         _sentryAIOptions = new SentryAIOptions();
         configure?.Invoke(_sentryAIOptions);
     }
@@ -102,9 +111,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
 
     /// <inheritdoc/>
     public override object? GetService(Type serviceType, object? serviceKey = null) =>
-        serviceType == typeof(ActivitySource)
-            ? SentryAIActivitySource.Instance
-            : base.GetService(serviceType, serviceKey);
+        serviceType == typeof(ActivitySource) ? _activitySource : base.GetService(serviceType, serviceKey);
 
     private void AfterResponseCleanup(ISpan chatSpan, ISpan agentSpan, Exception? exception = null)
     {

src/Sentry.Extensions.AI/SentryChatClient.cs

@@ -16,28 +16,25 @@ public Fixture()
 
     private readonly Fixture _fixture = new();
 
-    public SentryAIActivityListenerTests()
-    {
-        // Dispose ActivityListener before each test, otherwise the singleton instance will persist between tests
-        new SentryAiActivityListener().Dispose();
-    }
-
     [Fact]
     public void Init_AddsActivityListenerToActivitySource()
     {
+        // Arrange
+        var source = SentryAIActivitySource.CreateSource();
+
         // Act
-        _ = new SentryAiActivityListener();
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
 
         // Assert
-        Assert.True(SentryAIActivitySource.Instance.HasListeners());
+        Assert.True(source.HasListeners());
     }
 
-    [Theory]
-    [InlineData(SentryAIConstants.SentryActivitySourceName)]
-    public void ShouldListenTo_ReturnsTrueForSentryActivitySource(string sourceName)
+    [Fact]
+    public void ShouldListenTo_ReturnsTrueForSentryActivitySource()
     {
         // Arrange
-        _ = new SentryAiActivityListener(_fixture.Hub);
+        var sourceName = SentryAIActivitySource.SentryActivitySourceName;
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
         var activitySource = new ActivitySource(sourceName);
 
         // Act
@@ -57,7 +54,7 @@ public void ShouldListenTo_ReturnsTrueForSentryActivitySource(string sourceName)
     public void ShouldListenTo_ReturnsFalseForNonSentryActivitySource()
     {
         // Arrange
-        _ = new SentryAiActivityListener(_fixture.Hub);
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
         var activitySource = new ActivitySource("Other.ActivitySource");
 
         // Act & Assert
@@ -69,15 +66,16 @@ public void ShouldListenTo_ReturnsFalseForNonSentryActivitySource()
             Arg.Any<IReadOnlyDictionary<string, object?>>());
     }
 
-    [Theory]
-    [InlineData("orchestrate_tools")]
-    public void Sample_ReturnsAllDataAndRecordedForFICCActivityNames(string activityName)
+    [Fact]
+    public void Sample_ReturnsAllDataAndRecordedForFICCActivityNames()
     {
         // Arrange
-        _ = new SentryAiActivityListener(_fixture.Hub);
+        var activityName = "orchestrate_tools";
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
+        var source = SentryAIActivitySource.CreateSource();
 
         // Act
-        using var activity = SentryAIActivitySource.Instance.StartActivity(activityName);
+        using var activity = source.StartActivity(activityName);
 
         // Assert
         Assert.NotNull(activity);
@@ -89,25 +87,4 @@ public void Sample_ReturnsAllDataAndRecordedForFICCActivityNames(string activity
             Arg.Any<ITransactionContext>(),
             Arg.Any<IReadOnlyDictionary<string, object?>>());
     }
-
-    [Fact]
-    public void Init_MultipleCalls_NoDuplicateListener_StartsOnlyOneTransaction()
-    {
-        // Arrange
-        _ = new SentryAiActivityListener(_fixture.Hub);
-        _ = new SentryAiActivityListener(_fixture.Hub);
-        _ = new SentryAiActivityListener(_fixture.Hub);
-
-        // Act
-        using var activity = SentryAIActivitySource.Instance.StartActivity(SentryAIConstants.FICCActivityNames[0]);
-
-        // Assert
-        Assert.NotNull(activity);
-        Assert.True(SentryAIActivitySource.Instance.HasListeners());
-        activity.Stop();
-
-        _fixture.Hub.Received(1).StartTransaction(
-            Arg.Any<ITransactionContext>(),
-            Arg.Any<IReadOnlyDictionary<string, object?>>());
-    }
 }

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

@@ -5,14 +5,27 @@ namespace Sentry.Extensions.AI.Tests;
 
 public class SentryAIExtensionsTests
 {
+    private class Fixture
+    {
+        public IHub Hub { get; } = Substitute.For<IHub>();
+
+        public Fixture()
+        {
+            Hub.IsEnabled.Returns(true);
+        }
+    }
+
+    private readonly Fixture _fixture = new();
+
     [Fact]
     public void WithSentry_IChatClient_ReturnsWrappedClient()
     {
         // Arrange
         var mockClient = Substitute.For<IChatClient>();
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
 
         // Act
-        var result = mockClient.AddSentry();
+        var result = mockClient.AddSentry(listener);
 
         // Assert
         Assert.IsType<SentryChatClient>(result);
@@ -23,10 +36,11 @@ public void WithSentry_IChatClient_WithConfiguration_PassesConfigurationToWrappe
     {
         // Arrange
         var mockClient = Substitute.For<IChatClient>();
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
         var configureWasCalled = false;
 
         // Act
-        var result = mockClient.AddSentry(options =>
+        var result = mockClient.AddSentry(listener, options =>
             {
                 configureWasCalled = true;
                 options.Experimental.RecordInputs = false;
@@ -44,9 +58,10 @@ public void WithSentry_IChatClient_WithNullConfiguration_UsesDefaultConfiguratio
     {
         // Arrange
         var mockClient = Substitute.For<IChatClient>();
+        using var listener = SentryAiActivityListener.CreateListener(_fixture.Hub);
 
         // Act
-        var result = mockClient.AddSentry(null);
+        var result = mockClient.AddSentry(listener, null);
 
         // Assert
         Assert.IsType<SentryChatClient>(result);