fix(otel): set invocation span metadata

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py

@@ -7,7 +7,10 @@
 import threading
 from typing import Any
 
-from aws_durable_execution_sdk_python.lambda_service import OperationType
+from aws_durable_execution_sdk_python.lambda_service import (
+    InvocationStatus,
+    OperationType,
+)
 from aws_durable_execution_sdk_python.plugin import (
     DurableInstrumentationPlugin,
     InvocationEndInfo,
@@ -25,6 +28,7 @@
     Link,
     Span,
     SpanContext,
+    SpanKind,
     StatusCode,
     TraceFlags,
     Tracer,
@@ -43,6 +47,8 @@
 
 logger = logging.getLogger(__name__)
 
+_SpanAttributes = dict[str, str | bool | int]
+
 
 def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None:
     """Convert a datetime to OTel timestamp (nanoseconds since epoch), or None."""
@@ -211,7 +217,7 @@ def _start_span(
         self,
         operation_id: str | None,
         name: str,
-        attributes: dict[str, str],
+        attributes: _SpanAttributes,
         start_time: datetime.datetime | None = None,
         parent_span: Span | None = None,
         existed: bool = False,
@@ -282,6 +288,7 @@ def _start_span(
                 )
             span = self._tracer.start_span(
                 name=name,
+                kind=SpanKind.INTERNAL,
                 attributes=attributes,
                 start_time=_to_otel_timestamp(start_time),
                 context=parent_context,
@@ -338,6 +345,18 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
             if operation_id:
                 self._end_span(operation_id)
 
+        invocation_span = self._get_span(None)
+        if invocation_span:
+            invocation_span.set_attribute(
+                "durable.invocation.status", info.status.value
+            )
+            if info.status is InvocationStatus.FAILED:
+                invocation_span.set_status(
+                    StatusCode.ERROR, info.error.message if info.error else ""
+                )
+            elif info.status is InvocationStatus.SUCCEEDED:
+                invocation_span.set_status(StatusCode.OK)
+
         # end the invocation span
         self._end_span(None)
 
@@ -512,7 +531,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
                 trace.set_span_in_context(parent_span, self._extracted_context)
             )
 
-    def _extract_attributes(self, info: Any) -> dict[str, str]:
+    def _extract_attributes(self, info: Any) -> _SpanAttributes:
         """Extract durable execution fields as OpenTelemetry span attributes.
 
         Args:
@@ -521,10 +540,12 @@ def _extract_attributes(self, info: Any) -> dict[str, str]:
         Returns:
             A dictionary of durable execution attributes suitable for a span.
         """
-        attributes: dict[str, str] = {
+        attributes: _SpanAttributes = {
             "durable.execution.arn": self._execution_arn,
         }
 
+        if isinstance(info, InvocationStartInfo):
+            attributes["durable.invocation.first"] = info.is_first_invocation
         if hasattr(info, "operation_id") and info.operation_id is not None:
             attributes["durable.operation.id"] = info.operation_id
         if hasattr(info, "operation_type") and info.operation_type is not None:

packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py

@@ -27,6 +27,7 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import SpanKind, StatusCode
 
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     operation_id_to_span_id,
@@ -65,24 +66,28 @@ def _create_plugin() -> tuple[OtelPlugin, InMemorySpanExporter]:
     return plugin, exporter
 
 
-def _invocation_start_info() -> InvocationStartInfo:
+def _invocation_start_info(
+    is_first_invocation: bool = True,
+) -> InvocationStartInfo:
     """Create standard invocation start info for tests."""
     return InvocationStartInfo(
         request_id="request-1",
         execution_arn=EXECUTION_ARN,
         execution_start_time=START_TIME,
-        is_first_invocation=True,
+        is_first_invocation=is_first_invocation,
     )
 
 
-def _invocation_end_info() -> InvocationEndInfo:
+def _invocation_end_info(
+    status: InvocationStatus = InvocationStatus.SUCCEEDED,
+) -> InvocationEndInfo:
     """Create standard invocation end info for tests."""
     return InvocationEndInfo(
         request_id="request-1",
         execution_arn=EXECUTION_ARN,
         execution_start_time=START_TIME,
         is_first_invocation=True,
-        status=InvocationStatus.SUCCEEDED,
+        status=status,
         error=None,
     )
 
@@ -144,10 +149,55 @@ def test_invocation_start_and_end_emit_invocation_span():
 
     spans = exporter.get_finished_spans()
     assert [span.name for span in spans] == ["invocation"]
+    assert spans[0].kind is SpanKind.INTERNAL
     assert spans[0].attributes["durable.execution.arn"] == EXECUTION_ARN
+    assert spans[0].attributes["durable.invocation.first"] is True
+    assert (
+        spans[0].attributes["durable.invocation.status"]
+        == InvocationStatus.SUCCEEDED.value
+    )
     assert plugin._get_span(None) is None
 
 
+def test_invocation_span_records_subsequent_invocation():
+    """Invocation spans preserve a false first-invocation attribute."""
+    plugin, exporter = _create_plugin()
+
+    plugin.on_invocation_start(_invocation_start_info(is_first_invocation=False))
+    plugin.on_invocation_end(_invocation_end_info())
+
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].attributes["durable.invocation.first"] is False
+
+
+@pytest.mark.parametrize(
+    ("invocation_status", "expected_span_status"),
+    [
+        (InvocationStatus.PENDING, StatusCode.UNSET),
+        (InvocationStatus.RETRY, StatusCode.UNSET),
+        (InvocationStatus.SUCCEEDED, StatusCode.OK),
+        (InvocationStatus.FAILED, StatusCode.ERROR),
+    ],
+)
+def test_invocation_span_status_reflects_execution_status(
+    invocation_status: InvocationStatus,
+    expected_span_status: StatusCode,
+):
+    """Only terminal invocation spans receive a success or failure status."""
+    plugin, exporter = _create_plugin()
+
+    plugin.on_invocation_start(_invocation_start_info())
+    plugin.on_invocation_end(_invocation_end_info(invocation_status))
+
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1
+    attributes = spans[0].attributes
+    assert attributes is not None
+    assert attributes["durable.invocation.status"] == invocation_status.value
+    assert spans[0].status.status_code is expected_span_status
+
+
 def test_operation_callbacks_emit_child_span_with_deterministic_span_id():
     """Verify non-user-function operations are traced beneath the invocation."""
     plugin, exporter = _create_plugin()
@@ -193,6 +243,7 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id():
     plugin.on_invocation_end(_invocation_end_info())
 
     spans_by_name = {span.name: span for span in exporter.get_finished_spans()}
+    assert all(span.kind is SpanKind.INTERNAL for span in spans_by_name.values())
     wait_span = spans_by_name["wait-for-signal"]
     invocation_span = spans_by_name["invocation"]
     assert wait_span.context.span_id == operation_id_to_span_id(
@@ -444,6 +495,8 @@ def test_step_operation_span_parents_attempt_span():
     spans_by_name = {span.name: span for span in exporter.get_finished_spans()}
     finished_step_span = spans_by_name["fetch-user"]
     attempt_span = spans_by_name["fetch-user attempt 1"]
+    assert finished_step_span.kind is SpanKind.INTERNAL
+    assert attempt_span.kind is SpanKind.INTERNAL
     assert attempt_span.parent.span_id == finished_step_span.context.span_id
     assert (
         finished_step_span.attributes["durable.operation.status"]