Merge branch 'master' into webb/log-batcher-upper-bound

Author: alexander-alderman-webb

sentry_sdk/_metrics_batcher.py

@@ -13,15 +13,18 @@
 
 class MetricsBatcher:
     MAX_METRICS_BEFORE_FLUSH = 1000
+    MAX_METRICS_BEFORE_DROP = 10_000
     FLUSH_WAIT_TIME = 5.0
 
     def __init__(
         self,
         capture_func,  # type: Callable[[Envelope], None]
+        record_lost_func,  # type: Callable[..., None]
     ):
         # type: (...) -> None
         self._metric_buffer = []  # type: List[Metric]
         self._capture_func = capture_func
+        self._record_lost_func = record_lost_func
         self._running = True
         self._lock = threading.Lock()
 
@@ -72,6 +75,14 @@ def add(
             return None
 
         with self._lock:
+            if len(self._metric_buffer) >= self.MAX_METRICS_BEFORE_DROP:
+                self._record_lost_func(
+                    reason="queue_overflow",
+                    data_category="trace_metric",
+                    quantity=1,
+                )
+                return None
+
             self._metric_buffer.append(metric)
             if len(self._metric_buffer) >= self.MAX_METRICS_BEFORE_FLUSH:
                 self._flush_event.set()

sentry_sdk/client.py

@@ -393,7 +393,10 @@ def _record_lost_event(
 
             self.metrics_batcher = None
             if has_metrics_enabled(self.options):
-                self.metrics_batcher = MetricsBatcher(capture_func=_capture_envelope)
+                self.metrics_batcher = MetricsBatcher(
+                    capture_func=_capture_envelope,
+                    record_lost_func=_record_lost_event,
+                )
 
             max_request_body_size = ("always", "never", "small", "medium")
             if self.options["max_request_body_size"] not in max_request_body_size:

sentry_sdk/consts.py

@@ -1012,6 +1012,7 @@ def __init__(
         before_send_log=None,  # type: Optional[Callable[[Log, Hint], Optional[Log]]]
         trace_ignore_status_codes=frozenset(),  # type: AbstractSet[int]
         enable_metrics=True,  # type: bool
+        before_send_metric=None,  # type: Optional[Callable[[Metric, Hint], Optional[Metric]]]
     ):
         # type: (...) -> None
         """Initialize the Sentry SDK with the given parameters. All parameters described here can be used in a call to `sentry_sdk.init()`.

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -1,10 +1,9 @@
 from functools import wraps
 
 import sentry_sdk
-from sentry_sdk.tracing_utils import set_span_errored
-from sentry_sdk.utils import event_from_exception
 
 from ..spans import invoke_agent_span, update_invoke_agent_span
+from ..utils import _capture_exception
 
 from typing import TYPE_CHECKING
 from pydantic_ai.agent import Agent  # type: ignore
@@ -13,18 +12,6 @@
     from typing import Any, Callable, Optional
 
 
-def _capture_exception(exc):
-    # type: (Any) -> None
-    set_span_errored()
-
-    event, hint = event_from_exception(
-        exc,
-        client_options=sentry_sdk.get_client().options,
-        mechanism={"type": "pydantic_ai", "handled": False},
-    )
-    sentry_sdk.capture_event(event, hint=hint)
-
-
 class _StreamingContextManagerWrapper:
     """Wrapper for streaming methods that return async context managers."""
 

sentry_sdk/integrations/pydantic_ai/patches/tools.py

@@ -5,6 +5,7 @@
 import sentry_sdk
 
 from ..spans import execute_tool_span, update_execute_tool_span
+from ..utils import _capture_exception
 
 from typing import TYPE_CHECKING
 
@@ -56,16 +57,21 @@ async def wrapped_call_tool(self, call, allow_partial, wrap_validation_errors):
             )
             agent = agent_data.get("_agent")
 
-            # Get args for span (before validation)
-            # call.args can be a string (JSON) or dict
-            args_dict = call.args if isinstance(call.args, dict) else {}
+            try:
+                args_dict = call.args_as_dict()
+            except Exception:
+                args_dict = call.args if isinstance(call.args, dict) else {}
 
             with execute_tool_span(name, args_dict, agent, tool_type=tool_type) as span:
-                result = await original_call_tool(
-                    self, call, allow_partial, wrap_validation_errors
-                )
-                update_execute_tool_span(span, result)
-                return result
+                try:
+                    result = await original_call_tool(
+                        self, call, allow_partial, wrap_validation_errors
+                    )
+                    update_execute_tool_span(span, result)
+                    return result
+                except Exception as exc:
+                    _capture_exception(exc)
+                    raise exc from None
 
         # No span context - just call original
         return await original_call_tool(

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -1,14 +1,13 @@
 import sentry_sdk
-from sentry_sdk.ai.utils import set_data_normalized
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.utils import safe_serialize
+from sentry_sdk.tracing_utils import set_span_errored
+from sentry_sdk.utils import event_from_exception, safe_serialize
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, List, Dict
-    from pydantic_ai.usage import RequestUsage  # type: ignore
+    from typing import Any
 
 
 def _should_send_prompts():
@@ -173,3 +172,15 @@ def _set_available_tools(span, agent):
     except Exception:
         # If we can't extract tools, just skip it
         pass
+
+
+def _capture_exception(exc):
+    # type: (Any) -> None
+    set_span_errored()
+
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": "pydantic_ai", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)

sentry_sdk/utils.py

@@ -2060,4 +2060,6 @@ def get_before_send_metric(options):
     if options is None:
         return None
 
-    return options["_experiments"].get("before_send_metric")
+    return options.get("before_send_metric") or options["_experiments"].get(
+        "before_send_metric"
+    )

tests/test_metrics.py

@@ -168,6 +168,45 @@ def test_metrics_tracing_without_performance(sentry_init, capture_envelopes):
 def test_metrics_before_send(sentry_init, capture_envelopes):
     before_metric_called = False
 
+    def _before_metric(record, hint):
+        nonlocal before_metric_called
+
+        assert set(record.keys()) == {
+            "timestamp",
+            "trace_id",
+            "span_id",
+            "name",
+            "type",
+            "value",
+            "unit",
+            "attributes",
+        }
+
+        if record["name"] == "test.skip":
+            return None
+
+        before_metric_called = True
+        return record
+
+    sentry_init(
+        before_send_metric=_before_metric,
+    )
+    envelopes = capture_envelopes()
+
+    sentry_sdk.metrics.count("test.skip", 1)
+    sentry_sdk.metrics.count("test.keep", 1)
+
+    get_client().flush()
+
+    metrics = envelopes_to_metrics(envelopes)
+    assert len(metrics) == 1
+    assert metrics[0]["name"] == "test.keep"
+    assert before_metric_called
+
+
+def test_metrics_experimental_before_send(sentry_init, capture_envelopes):
+    before_metric_called = False
+
     def _before_metric(record, hint):
         nonlocal before_metric_called
 
@@ -204,3 +243,27 @@ def _before_metric(record, hint):
     assert len(metrics) == 1
     assert metrics[0]["name"] == "test.keep"
     assert before_metric_called
+
+
+def test_batcher_drops_metrics(sentry_init, monkeypatch):
+    sentry_init()
+    client = sentry_sdk.get_client()
+
+    def no_op_flush():
+        pass
+
+    monkeypatch.setattr(client.metrics_batcher, "_flush", no_op_flush)
+
+    lost_event_calls = []
+
+    def record_lost_event(reason, data_category, quantity):
+        lost_event_calls.append((reason, data_category, quantity))
+
+    monkeypatch.setattr(client.metrics_batcher, "_record_lost_func", record_lost_event)
+
+    for i in range(10_005):  # 5 metrics over the hard limit
+        sentry_sdk.metrics.count("test.counter", 1)
+
+    assert len(lost_event_calls) == 5
+    for lost_event_call in lost_event_calls:
+        assert lost_event_call == ("queue_overflow", "trace_metric", 1)