From 329526a1657218b1729e857b945019c64c5a33cc Mon Sep 17 00:00:00 2001
From: Giselle van Dongen <giselle@restate.dev>
Date: Tue, 7 Apr 2026 14:06:28 +0200
Subject: [PATCH 1/2] Logfire Tracing: fix default arguments for Tracer

---
 python/restate/ext/tracing/_tracing.py | 79 ++++++++++++++------------
 1 file changed, 44 insertions(+), 35 deletions(-)

diff --git a/python/restate/ext/tracing/_tracing.py b/python/restate/ext/tracing/_tracing.py
index b272185..b4dc9c9 100644
--- a/python/restate/ext/tracing/_tracing.py
+++ b/python/restate/ext/tracing/_tracing.py
@@ -9,8 +9,13 @@
     # All spans created by this tracer are flat children of the Restate trace.
 """
 
-from opentelemetry.trace import INVALID_SPAN, use_span, Tracer, TracerProvider
+from typing import Optional, Iterator
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import INVALID_SPAN, Span, SpanKind, Tracer, TracerProvider, use_span, _Links
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.util import types
+from opentelemetry.util._decorator import _agnosticcontextmanager
 from restate.server_context import (
     current_context,
     get_extension_data,
@@ -58,15 +63,15 @@ def _get_root_context():
 
     def start_span(
         self,
-        name,
-        context=None,
-        kind=None,
-        attributes=None,
-        links=None,
-        start_time=None,
-        record_exception=True,
-        set_status_on_exception=True,
-    ):
+        name: str,
+        context: Optional[context_api.Context] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: types.Attributes = None,
+        links: _Links = None,
+        start_time: Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+    ) -> Span:
         if restate_context_is_replaying.get(False):
             return INVALID_SPAN
         root = self._get_root_context()
@@ -85,34 +90,38 @@ def start_span(
         self._track_span(span)
         return span
 
+    @_agnosticcontextmanager
     def start_as_current_span(
         self,
-        name,
-        context=None,
-        kind=None,
-        attributes=None,
-        links=None,
-        start_time=None,
-        record_exception=True,
-        set_status_on_exception=True,
-        end_on_exit=True,
-    ):
+        name: str,
+        context: Optional[context_api.Context] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: types.Attributes = None,
+        links: _Links = None,
+        start_time: Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+        end_on_exit: bool = True,
+    ) -> Iterator[Span]:
         if restate_context_is_replaying.get(False):
-            return use_span(INVALID_SPAN, end_on_exit=False)
-        root = self._get_root_context()
-        if root is not None:
-            context = root
-        return self._tracer.start_as_current_span(
-            name,
-            context=context,
-            kind=kind,
-            attributes=attributes,
-            links=links,
-            start_time=start_time,
-            record_exception=record_exception,
-            set_status_on_exception=set_status_on_exception,
-            end_on_exit=end_on_exit,
-        )
+            with use_span(INVALID_SPAN, end_on_exit=False) as span:
+                yield span
+        else:
+            root = self._get_root_context()
+            if root is not None:
+                context = root
+            with self._tracer.start_as_current_span(
+                name,
+                context=context,
+                kind=kind,
+                attributes=attributes,
+                links=links,
+                start_time=start_time,
+                record_exception=record_exception,
+                set_status_on_exception=set_status_on_exception,
+                end_on_exit=end_on_exit,
+            ) as span:
+                yield span
 
     @staticmethod
     def _track_span(span):

From 695f069481fe1253b3c693b63a11d2bbefb19ede Mon Sep 17 00:00:00 2001
From: Giselle van Dongen <giselle@restate.dev>
Date: Tue, 7 Apr 2026 15:01:29 +0200
Subject: [PATCH 2/2] Logfire Tracing: fix default arguments for Tracer

---
 python/restate/ext/tracing/_tracing.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/restate/ext/tracing/_tracing.py b/python/restate/ext/tracing/_tracing.py
index b4dc9c9..198fc93 100644
--- a/python/restate/ext/tracing/_tracing.py
+++ b/python/restate/ext/tracing/_tracing.py
@@ -9,10 +9,10 @@
     # All spans created by this tracer are flat children of the Restate trace.
 """
 
-from typing import Optional, Iterator
+from typing import Optional, Iterator, Sequence
 
 from opentelemetry import context as context_api
-from opentelemetry.trace import INVALID_SPAN, Span, SpanKind, Tracer, TracerProvider, use_span, _Links
+from opentelemetry.trace import INVALID_SPAN, Span, SpanKind, Tracer, TracerProvider, use_span, Link
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from opentelemetry.util import types
 from opentelemetry.util._decorator import _agnosticcontextmanager
@@ -67,7 +67,7 @@ def start_span(
         context: Optional[context_api.Context] = None,
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
-        links: _Links = None,
+        links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
@@ -97,7 +97,7 @@ def start_as_current_span(
         context: Optional[context_api.Context] = None,
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
-        links: _Links = None,
+        links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,