WIP feat(AI): Microsoft.Extensions.AI instrumentation
#4657
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #4657 +/- ##
==========================================
+ Coverage 73.49% 73.60% +0.11%
==========================================
Files 483 488 +5
Lines 17692 17843 +151
Branches 3492 3531 +39
==========================================
+ Hits 13002 13134 +132
- Misses 3799 3809 +10
- Partials 891 900 +9 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@sentry review |
![seer-by-sentry[bot]](https://avatars.githubusercontent.com/in/801464?s=60&v=4){kind=small}
|
@sentry review |
|
@BugBot review |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
![sentry-io[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
|
@BugBot review |
|
@sentry review |
| { | ||
| private readonly HubAdapter _hub; | ||
| private readonly SentryAIOptions _sentryAIOptions; | ||
| internal static ISpan? RootSpan; |
There was a problem hiding this comment.
Bug: Async Methods Share Unprotected Static Field
The static field RootSpan is accessed and modified from multiple async methods (GetResponseAsync and GetStreamingResponseAsync) without synchronization. This creates a race condition where concurrent calls can incorrectly share or overwrite the root span. When multiple concurrent AI chat operations occur, they may interfere with each other's span tracking, causing incorrect span hierarchies and data corruption.
There was a problem hiding this comment.
Hmm, I just tried wrapping RootSpan with AsyncLocal, and it resulted in something like this:
Compare it to the previous AI view we have which used normal ISpan:
A couple of problems:
-
The tool spans are not marked as a child of the
invoke_agentroot span, because tool calls are invoked in a separate await call inFunctionInvokingChatClient(fromMicrosoft.Extensions.AINOTAbstractions): link. -
The outer spans (
invoke_agentspans) are generated each time the LLM has to create a response after the tool call has finished.
I definitely do agree that 2 concurrent calls to GetResponseAsync could result in a wonky situation. I'll come up with a test to reproduce this issue and try to come up with a solution...
There was a problem hiding this comment.
The transaction created by a call like this will be asynclocal:
See #3596
So I don't think there is any other option that won't break things.
| { | ||
| private readonly HubAdapter _hub; | ||
| private readonly SentryAIOptions _sentryAIOptions; | ||
| internal static ISpan? RootSpan; |
There was a problem hiding this comment.
Bug: RootSpan is a static field, not AsyncLocal, causing incorrect span hierarchy and potential NullReferenceException in concurrent operations.
Severity: HIGH | Confidence: 0.95
🔍 Detailed Analysis
The RootSpan field is declared as internal static ISpan? RootSpan; which is a static field. This design choice violates the established pattern within the codebase, where AsyncLocal<T> is used for transactions and spans to prevent cross-thread contamination. In concurrent execution, one thread might set RootSpan = null while another thread is still attempting to use it, leading to an incorrect span hierarchy or a NullReferenceException when accessing the field.
💡 Suggested Fix
Change internal static ISpan? RootSpan; to internal static AsyncLocal<ISpan?> RootSpan; and adjust its usage to correctly manage span context across asynchronous operations.
🤖 Prompt for AI Agent
Fix this bug. In src/Sentry.Extensions.AI/SentryChatClient.cs at line 10: The `RootSpan`
field is declared as `internal static ISpan? RootSpan;` which is a static field. This
design choice violates the established pattern within the codebase, where
`AsyncLocal<T>` is used for transactions and spans to prevent cross-thread
contamination. In concurrent execution, one thread might set `RootSpan = null` while
another thread is still attempting to use it, leading to an incorrect span hierarchy or
a `NullReferenceException` when accessing the field.
Did we get this right? 👍 / 👎 to inform future reviews.
| internal sealed class SentryChatClient : DelegatingChatClient | ||
| { | ||
| private readonly HubAdapter _hub; | ||
| private readonly SentryAIOptions _sentryAIOptions; |
There was a problem hiding this comment.
CRITICAL: The static RootSpan field is not thread-safe and will cause issues in concurrent scenarios. When multiple async operations execute concurrently, they all share the same static field, leading to race conditions where one thread's operation can interfere with another's. This should use AsyncLocal<ISpan?> instead to provide proper async context isolation.
Example problematic scenario:
- Thread A starts GetResponseAsync, creates RootSpan
- Thread B starts GetResponseAsync, reads the same RootSpan
- Thread A completes and sets RootSpan = null
- Thread B tries to use RootSpan but it's null or stale
Consider using AsyncLocal<ISpan?> or storing the span in the HubAdapter's scope.
Severity: CRITICAL
🤖 Prompt for AI Agent
Fix this code. In src/Sentry.Extensions.AI/SentryChatClient.cs#L9: CRITICAL: The static
`RootSpan` field is not thread-safe and will cause issues in concurrent scenarios. When
multiple async operations execute concurrently, they all share the same static field,
leading to race conditions where one thread's operation can interfere with another's.
This should use `AsyncLocal<ISpan?>` instead to provide proper async context isolation.
Example problematic scenario:
1. Thread A starts GetResponseAsync, creates RootSpan
2. Thread B starts GetResponseAsync, reads the same RootSpan
3. Thread A completes and sets RootSpan = null
4. Thread B tries to use RootSpan but it's null or stale
Consider using AsyncLocal<ISpan?> or storing the span in the HubAdapter's scope.
Did we get this right? 👍 / 👎 to inform future reviews.
| var outerSpan = EnsureRootSpanExists(); | ||
| var innerSpan = CreateChatSpan(options, outerSpan); | ||
|
|
||
| try | ||
| { | ||
| var chatMessages = messages as ChatMessage[] ?? messages.ToArray(); | ||
| SentryAISpanEnricher.EnrichWithRequest(innerSpan, chatMessages, options, _sentryAIOptions); | ||
|
|
||
| var response = await base.GetResponseAsync(chatMessages, options, cancellationToken).ConfigureAwait(false); | ||
|
|
||
| SentryAISpanEnricher.EnrichWithResponse(innerSpan, response, _sentryAIOptions); | ||
| innerSpan.Finish(SpanStatus.Ok); | ||
|
|
||
| // Only finish the outerSpan (RootSpan) if this response's finish reason is stop (not tool calls). | ||
| // This allows the RootSpan to persist throughout multiple `GetResponseAsync` calls | ||
| // happening before and after tool calls | ||
| var isResponseTerminal = response.FinishReason == ChatFinishReason.Stop; | ||
| if (isResponseTerminal) | ||
| { | ||
| outerSpan.Finish(SpanStatus.Ok); | ||
| RootSpan = null; | ||
| } | ||
|
|
||
| return response; | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| innerSpan.Finish(ex); | ||
| outerSpan.Finish(ex); | ||
| _hub.CaptureException(ex); | ||
| RootSpan = null; | ||
| throw; | ||
| } | ||
| } | ||
|
|
||
| /// <inheritdoc cref="IChatClient"/> | ||
| public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, | ||
| ChatOptions? options = null, | ||
| [EnumeratorCancellation] CancellationToken cancellationToken = new()) | ||
| { | ||
| var outerSpan = EnsureRootSpanExists(); | ||
| var innerSpan = CreateChatSpan(options, outerSpan); | ||
|
|
||
| var responses = new List<ChatResponseUpdate>(); | ||
| var chatMessages = messages as ChatMessage[] ?? messages.ToArray(); | ||
| var enumerator = base | ||
| .GetStreamingResponseAsync(chatMessages, options, cancellationToken) | ||
| .GetAsyncEnumerator(cancellationToken); | ||
| SentryAISpanEnricher.EnrichWithRequest(innerSpan, chatMessages, options, _sentryAIOptions); | ||
| ChatResponseUpdate? current = null; | ||
|
|
||
| while (true) | ||
| { | ||
| try | ||
| { | ||
| var hasNext = await enumerator.MoveNextAsync().ConfigureAwait(false); | ||
| if (!hasNext) | ||
| { | ||
| SentryAISpanEnricher.EnrichWithStreamingResponse(innerSpan, responses, _sentryAIOptions); | ||
| innerSpan.Finish(SpanStatus.Ok); | ||
|
|
||
| // Only if currentFinishReason is to stop, then we finish the RootSpan and set it to null. | ||
| // This allows the RootSpan to persist throughout multiple `GetStreamingResponseAsync` calls | ||
| // happening before and after tool calls | ||
| var shouldFinishRootSpan = current?.FinishReason == ChatFinishReason.Stop; | ||
| if (shouldFinishRootSpan) | ||
| { | ||
| outerSpan.Finish(SpanStatus.Ok); | ||
| RootSpan = null; | ||
| } | ||
|
|
||
| yield break; | ||
| } | ||
|
|
||
| current = enumerator.Current; | ||
| responses.Add(enumerator.Current); | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| innerSpan.Finish(ex); | ||
| outerSpan.Finish(ex); | ||
| _hub.CaptureException(ex); | ||
| RootSpan = null; | ||
| throw; | ||
| } | ||
|
|
||
| yield return current; | ||
| } | ||
| } | ||
|
|
||
| private ISpan EnsureRootSpanExists() | ||
| { | ||
| if (RootSpan == null) |
There was a problem hiding this comment.
CRITICAL: All direct accesses to the static RootSpan field (lines 30, 50, 70, 102, 122) need to be updated to use AsyncLocal<T> instead. This is necessary to fix the thread-safety issue. Ensure all reads and writes go through the AsyncLocal property.
Severity: CRITICAL
🤖 Prompt for AI Agent
Fix this code. In src/Sentry.Extensions.AI/SentryChatClient.cs#L30-L122: CRITICAL: All
direct accesses to the static `RootSpan` field (lines 30, 50, 70, 102, 122) need to be
updated to use AsyncLocal<T> instead. This is necessary to fix the thread-safety issue.
Ensure all reads and writes go through the AsyncLocal property.
Did we get this right? 👍 / 👎 to inform future reviews.
| using Microsoft.Extensions.AI; | ||
|
|
||
| namespace Sentry.Extensions.AI; |
There was a problem hiding this comment.
Missing required using statements. The class uses JsonSerializer (line 152) and StringBuilder (line 119) but doesn't have explicit using directives for System.Text.Json and System.Text. Add these using statements at the top of the file to make dependencies explicit and avoid relying on transitive using statements.
Severity: MEDIUM
🤖 Prompt for AI Agent
Fix this code. In src/Sentry.Extensions.AI/SentryAISpanEnricher.cs#L1-L3: Missing
required using statements. The class uses `JsonSerializer` (line 152) and
`StringBuilder` (line 119) but doesn't have explicit using directives for
`System.Text.Json` and `System.Text`. Add these using statements at the top of the file
to make dependencies explicit and avoid relying on transitive using statements.
Did we get this right? 👍 / 👎 to inform future reviews.
| var spanName = $"execute_tool {Name}"; | ||
| var currSpan = SentryChatClient.RootSpan == null ? | ||
| _hub.StartSpan(operation, spanName) : |
There was a problem hiding this comment.
This class directly accesses the static SentryChatClient.RootSpan field (line 16). Once you refactor RootSpan to use AsyncLocal<T> in SentryChatClient, you'll need to change this access pattern. Consider adding a public property or method to SentryChatClient to access the current span in a thread-safe manner rather than accessing the internal static field directly.
Severity: HIGH
🤖 Prompt for AI Agent
Fix this code. In src/Sentry.Extensions.AI/SentryInstrumentedFunction.cs#L16-L18: This
class directly accesses the static `SentryChatClient.RootSpan` field (line 16). Once you
refactor RootSpan to use AsyncLocal<T> in SentryChatClient, you'll need to change this
access pattern. Consider adding a public property or method to SentryChatClient to
access the current span in a thread-safe manner rather than accessing the internal
static field directly.
Did we get this right? 👍 / 👎 to inform future reviews.
| // happening before and after tool calls | ||
| var isResponseTerminal = response.FinishReason == ChatFinishReason.Stop; | ||
| if (isResponseTerminal) | ||
| { | ||
| outerSpan.Finish(SpanStatus.Ok); | ||
| RootSpan = null; | ||
| } | ||
|
|
||
| return response; | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| innerSpan.Finish(ex); | ||
| outerSpan.Finish(ex); | ||
| _hub.CaptureException(ex); | ||
| RootSpan = null; | ||
| throw; | ||
| } | ||
| } | ||
|
|
||
| /// <inheritdoc cref="IChatClient"/> | ||
| public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, | ||
| ChatOptions? options = null, | ||
| [EnumeratorCancellation] CancellationToken cancellationToken = new()) | ||
| { | ||
| var outerSpan = EnsureRootSpanExists(); | ||
| var innerSpan = CreateChatSpan(options, outerSpan); | ||
|
|
||
| var responses = new List<ChatResponseUpdate>(); | ||
| var chatMessages = messages as ChatMessage[] ?? messages.ToArray(); | ||
| var enumerator = base | ||
| .GetStreamingResponseAsync(chatMessages, options, cancellationToken) | ||
| .GetAsyncEnumerator(cancellationToken); | ||
| SentryAISpanEnricher.EnrichWithRequest(innerSpan, chatMessages, options, _sentryAIOptions); | ||
| ChatResponseUpdate? current = null; | ||
|
|
||
| while (true) | ||
| { | ||
| try | ||
| { | ||
| var hasNext = await enumerator.MoveNextAsync().ConfigureAwait(false); | ||
| if (!hasNext) | ||
| { | ||
| SentryAISpanEnricher.EnrichWithStreamingResponse(innerSpan, responses, _sentryAIOptions); | ||
| innerSpan.Finish(SpanStatus.Ok); | ||
|
|
||
| // Only if currentFinishReason is to stop, then we finish the RootSpan and set it to null. | ||
| // This allows the RootSpan to persist throughout multiple `GetStreamingResponseAsync` calls | ||
| // happening before and after tool calls | ||
| var shouldFinishRootSpan = current?.FinishReason == ChatFinishReason.Stop; | ||
| if (shouldFinishRootSpan) | ||
| { | ||
| outerSpan.Finish(SpanStatus.Ok); | ||
| RootSpan = null; | ||
| } | ||
|
|
||
| yield break; | ||
| } | ||
|
|
||
| current = enumerator.Current; | ||
| responses.Add(enumerator.Current); | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| innerSpan.Finish(ex); | ||
| outerSpan.Finish(ex); |
There was a problem hiding this comment.
The logic for managing RootSpan cleanup is complex and duplicated across both GetResponseAsync and GetStreamingResponseAsync. The cleanup happens in two places (lines 50, 56, 100, 106). Consider extracting the RootSpan lifecycle management into a helper method to reduce duplication and ensure consistent behavior.
Severity: MEDIUM
🤖 Prompt for AI Agent
Fix this code. In src/Sentry.Extensions.AI/SentryChatClient.cs#L45-L110: The logic for
managing RootSpan cleanup is complex and duplicated across both GetResponseAsync and
GetStreamingResponseAsync. The cleanup happens in two places (lines 50, 56, 100, 106).
Consider extracting the RootSpan lifecycle management into a helper method to reduce
duplication and ensure consistent behavior.
Did we get this right? 👍 / 👎 to inform future reviews.
| #nullable enable | ||
| using Microsoft.Extensions.AI; | ||
| using Sentry.Extensions.AI; | ||
|
|
||
| namespace Sentry.Extensions.AI.Tests; | ||
|
|
||
| public class SentryChatClientTests | ||
| { | ||
| [Fact] | ||
| public async Task CompleteAsync_CallsInnerClient() | ||
| { | ||
| var inner = Substitute.For<IChatClient>(); | ||
| var message = new ChatMessage(ChatRole.Assistant, "ok"); | ||
| var chatResponse = new ChatResponse(message); | ||
| inner.GetResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(), Arg.Any<CancellationToken>()) | ||
| .Returns(Task.FromResult(chatResponse)); | ||
|
|
||
| var sentryChatClient = new SentryChatClient(inner); | ||
|
|
||
| var res = await sentryChatClient.GetResponseAsync([new ChatMessage(ChatRole.User, "hi")], null); | ||
|
|
||
| Assert.Equal([message], res.Messages); | ||
| await inner.Received(1).GetResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(), Arg.Any<CancellationToken>()); | ||
| } | ||
|
|
||
| [Fact] | ||
| public async Task CompleteStreamingAsync_CallsInnerClient() | ||
| { | ||
| var inner = Substitute.For<IChatClient>(); | ||
|
|
||
| inner.GetStreamingResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(), Arg.Any<CancellationToken>()) | ||
| .Returns(CreateTestStreamingUpdatesAsync()); | ||
|
|
||
| var client = new SentryChatClient(inner); | ||
|
|
||
| var results = new List<ChatResponseUpdate>(); | ||
| await foreach (var update in client.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "hi")], null)) | ||
| { | ||
| results.Add(update); | ||
| } | ||
|
|
||
| Assert.Equal(2, results.Count); | ||
| Assert.Equal("Hello", results[0].Text); | ||
| Assert.Equal(" World!", results[1].Text); | ||
|
|
||
| inner.Received(1).GetStreamingResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(), Arg.Any<CancellationToken>()); | ||
| } |
There was a problem hiding this comment.
The test class doesn't properly clean up the static RootSpan state between tests. If tests run in parallel or if one test fails, it could leave RootSpan set and affect subsequent tests. Add proper cleanup in a test fixture or teardown method to reset the static state. This becomes even more important when the static field is converted to AsyncLocal<T>.
Severity: MEDIUM
🤖 Prompt for AI Agent
Fix this code. In test/Sentry.Extensions.AI.Tests/SentryChatClientTests.cs#L1-L47: The
test class doesn't properly clean up the static `RootSpan` state between tests. If tests
run in parallel or if one test fails, it could leave `RootSpan` set and affect
subsequent tests. Add proper cleanup in a test fixture or teardown method to reset the
static state. This becomes even more important when the static field is converted to
AsyncLocal<T>.
Did we get this right? 👍 / 👎 to inform future reviews.
3 participants
WIP!!
For now, here's a quick screenshot of what it looks like right now