Skip to content

feat(streaming): Support ai_track decorator #6895

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sentrivana merged 1 commit into from
Jul 27, 2026
+284 −51
Merged

feat(streaming): Support ai_track decorator #6895

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 128 additions & 50 deletions sentry_sdk/ai/monitoring.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from sentry_sdk.consts import SPANDATA
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing import Span
from sentry_sdk.tracing_utils import has_span_streaming_enabled
from sentry_sdk.utils import ContextVar, capture_internal_exceptions, reraise

if TYPE_CHECKING:
Expand All @@ -30,64 +31,141 @@ def get_ai_pipeline_name() -> "Optional[str]":
def ai_track(description: str, **span_kwargs: "Any") -> "Callable[[F], F]":
def decorator(f: "F") -> "F":
def sync_wrapped(*args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()

curr_pipeline = _ai_pipeline_name.get()
op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline")

@sentry sentry Bot Jul 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Using span_kwargs.pop("op") mutates the shared span_kwargs dictionary, causing subsequent calls to the decorated function to lose the custom op and use a default value.
Severity: MEDIUM

Suggested Fix

Instead of mutating span_kwargs with .pop(), use .get() to retrieve the op value. To prevent passing op along with other keyword arguments, create a copy of the dictionary for the start_span call, excluding the op key. For example: op = span_kwargs.get("op", ...) and other_kwargs = {k: v for k, v in span_kwargs.items() if k != 'op'}.

Prompt for AI Agent 
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: sentry_sdk/ai/monitoring.py#L37

Potential issue: The `span_kwargs` dictionary is created once when the `@ai_track`
decorator is applied and is shared across all calls to the decorated function. The code
uses `span_kwargs.pop("op", ...)` to retrieve the operation name, which mutates the
dictionary by permanently removing the `op` key. Consequently, only the first call to a
decorated function with a custom `op` (e.g., `@ai_track("description", op="custom.op")`)
will use the correct operation. All subsequent calls will find the `op` key missing and
will incorrectly fall back to the default value, such as "ai.run".

Did we get this right? 👍 / 👎 to inform future reviews.

@sentrivana sentrivana Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pre-existing issue


with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "ai_monitoring", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=description, attributes={"sentry.op": op}
) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_attribute(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_attribute(k, v)

if curr_pipeline:
Comment on lines +45 to +48

@sentry sentry Bot Jul 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: In streaming mode, extra keyword arguments passed to the @ai_track decorator, other than op, are silently ignored instead of being passed to start_span.
Severity: MEDIUM

Suggested Fix

Create a copy of the span_kwargs dictionary, remove the op key, and then pass the remaining arguments to the start_span call in the streaming path. This can be done with other_kwargs = {k: v for k, v in span_kwargs.items() if k != 'op'} and then spreading them into the function call: start_span(..., **other_kwargs).

Prompt for AI Agent 
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: sentry_sdk/ai/monitoring.py#L45-L48

Potential issue: In the new streaming mode code path for the `@ai_track` decorator, any
additional keyword arguments passed to the decorator (e.g., `origin`, `only_if_parent`)
are ignored. The code extracts the `op` but does not forward the rest of the
`span_kwargs` to the `sentry_sdk.traces.start_span` function. This creates an
inconsistency where the decorator's behavior changes depending on whether streaming is
active, silently dropping user-provided span keyword arguments in streaming mode while
correctly applying them in non-streaming mode.

Did we get this right? 👍 / 👎 to inform future reviews.

@sentrivana sentrivana Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct, intentional

span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={
"type": "ai_monitoring",
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res

else:
with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={
"type": "ai_monitoring",
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res

async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()

curr_pipeline = _ai_pipeline_name.get()
op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline")

with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return await f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = await f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "ai_monitoring", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=description, attributes={"sentry.op": op}
) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_attribute(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_attribute(k, v)

if curr_pipeline:
span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return await f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = await f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={
"type": "ai_monitoring",
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res
else:
with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline)
return await f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = await f(*args, **kwargs)
except Exception as e:
exc_info = sys.exc_info()
with capture_internal_exceptions():
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={
"type": "ai_monitoring",
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
reraise(*exc_info)
finally:
_ai_pipeline_name.set(None)
return res

if inspect.iscoroutinefunction(f):
return wraps(f)(async_wrapped) # type: ignore
Expand Down
Loading
Loading