fix(insight): preserve bounded renderer values

Author: Alex Wang

packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py

@@ -21,13 +21,16 @@
 from __future__ import annotations
 
 import copy
+import datetime
+import decimal
 import functools
 import itertools
 import logging
 import sys
 import threading
 import time
 import types
+import uuid
 from collections import OrderedDict, deque
 from typing import Any
 
@@ -69,6 +72,16 @@
 )
 
 
+_SAFE_OPAQUE_RETAINED_TYPES = (
+    datetime.date,
+    datetime.datetime,
+    datetime.time,
+    datetime.timedelta,
+    decimal.Decimal,
+    uuid.UUID,
+)
+
+
 class _RetainedChildren:
     __slots__ = ("iterator",)
 
@@ -113,29 +126,39 @@ def _retained_children(value: Any) -> Any:
     if isinstance(value, deque):
         return itertools.chain(deque.__iter__(value), custom)
     if isinstance(value, memoryview):
-        return (value.obj,)
+        return itertools.chain((value.obj,), custom)
     if isinstance(value, functools.partial):
-        return (value.func, value.args, value.keywords)
+        return itertools.chain((value.func, value.args, value.keywords), custom)
     if isinstance(value, types.FunctionType):
         closure = []
         for cell in value.__closure__ or ():
             try:
                 closure.append(cell.cell_contents)
             except ValueError:
                 pass
-        return itertools.chain(closure, (value.__defaults__, value.__kwdefaults__))
+        return itertools.chain(
+            closure, (value.__defaults__, value.__kwdefaults__), custom
+        )
     if isinstance(value, types.MethodType):
-        return (value.__self__, value.__func__)
+        return itertools.chain((value.__self__, value.__func__), custom)
     if isinstance(value, types.BuiltinFunctionType):
         owner = value.__self__
-        return () if owner is None or isinstance(owner, types.ModuleType) else (owner,)
+        retained = (
+            () if owner is None or isinstance(owner, types.ModuleType) else (owner,)
+        )
+        return itertools.chain(retained, custom)
     if isinstance(value, types.MethodWrapperType):
-        return (value.__self__,)
+        return itertools.chain((value.__self__,), custom)
     if isinstance(value, types.GeneratorType):
         frame = value.gi_frame
-        return () if frame is None else (frame.f_locals, value.gi_yieldfrom)
-    if isinstance(value, _ATOMIC_RETAINED_TYPES):
-        return ()
+        generator_children = (
+            () if frame is None else (frame.f_locals, value.gi_yieldfrom)
+        )
+        return itertools.chain(generator_children, custom)
+    if type(value) in _ATOMIC_RETAINED_TYPES:
+        return custom
+    if type(value) in _SAFE_OPAQUE_RETAINED_TYPES:
+        return custom
     if custom:
         return custom
     return _UNSUPPORTED_RETAINED_GRAPH
@@ -568,7 +591,16 @@ def __init__(
 
     def schedule(self, execution_arn: str, record: dict[str, Any]) -> None:
         """Fan a canonical record out to every lane. Returns immediately."""
-        record_size = _estimate_retained_size(record, self._max_pending_bytes)
+        try:
+            record_size = _estimate_retained_size(record, self._max_pending_bytes)
+        except Exception as exc:  # noqa: BLE001 - inspection must never break a hook
+            _logger.warning(
+                "workflow-insight: retained-size inspection failed for %s; "
+                "rejecting this record safely: %s",
+                execution_arn,
+                exc,
+            )
+            record_size = self._max_pending_bytes + 1
         for lane in self._lanes:
             lane.schedule(execution_arn, record, record_size)
 

packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py

@@ -11,12 +11,15 @@
 
 from __future__ import annotations
 
+import datetime
 import functools
 import threading
 import time
 import types
 from typing import Any
 
+import aws_durable_execution_sdk_python_insight._export_scheduler as scheduler_module
+
 from aws_durable_execution_sdk_python_insight._export_scheduler import (
     _ExportScheduler,
 )
@@ -721,12 +724,18 @@ def __iter__(self):
             self.iterated = True
             return iter(())
 
-    backing_buffers = [bytearray(4_000) for _ in range(3)]
+    class PayloadPartial(functools.partial):
+        pass
+
+    backing_buffers = [bytearray(4_000) for _ in range(4)]
     hidden = HiddenList(backing_buffers[2])
+    partial_with_payload = PayloadPartial(lambda value: value, "small")
+    partial_with_payload.payload = backing_buffers[3]
     payloads = [
         functools.partial(lambda value: value, backing_buffers[0]),
         (value for value in (backing_buffers[1],)),
         hidden,
+        partial_with_payload,
     ]
 
     for payload in payloads:
@@ -775,6 +784,42 @@ def closure() -> bytearray:
         scheduler.end_invocation(5.0)
 
 
+def test_safe_opaque_datetime_reaches_custom_renderer():
+    class DateRenderExporter(RecordingExporter):
+        def __init__(self) -> None:
+            super().__init__(max_record_size_bytes=10_000)
+            self.rendered: list[str] = []
+
+        def render(self, record: dict[str, Any]) -> Any:
+            value = record["payload"].isoformat()
+            self.rendered.append(value)
+            return {"value": value}
+
+    exporter = DateRenderExporter()
+    scheduler = _ExportScheduler([exporter])
+    record = _rec(ARN_A, "date")
+    record["payload"] = datetime.date(2026, 9, 10)
+    scheduler.schedule(ARN_A, record)
+    scheduler.end_invocation(5.0)
+
+    assert exporter.rendered == ["2026-09-10"]
+    assert exporter.exported_values() == ["date"]
+
+
+def test_retained_size_inspection_failure_does_not_escape_schedule(monkeypatch):
+    def fail_estimate(value: Any, max_size: int | None = None) -> int:
+        raise RuntimeError("inspection failed")
+
+    monkeypatch.setattr(scheduler_module, "_estimate_retained_size", fail_estimate)
+    exporter = RecordingExporter()
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500)
+
+    scheduler.schedule(ARN_A, _rec(ARN_A, "rejected"))
+
+    assert scheduler._lanes[0]._pending_count() == 0
+    assert scheduler._lanes[0]._worker is None
+
+
 def test_timed_out_barrier_flushes_eventually_and_worker_exits():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])