From 0969a103a0b16be84b7157b71e0673931fb87ee7 Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 09:12:41 +0200 Subject: [PATCH 1/6] Inject sampled traceparent on workflow execute calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Without a traceparent header, external workers inherit the API's HTTP span context via Temporal. In production the API's spans are frequently unsampled (ParentBasedTraceIdRatio), so workers produce no-op spans and traces never reach the collector — /trace/otel returns WF_1500. Adds a BeforeRequestHook that fires on any /execute path and injects a sampled W3C traceparent: forwarding the active OTEL span if it is already sampled, otherwise generating a fresh sampled one. An explicitly set traceparent header is never overwritten. --- src/mistralai/client/_hooks/registration.py | 2 ++ src/mistralai/client/_hooks/traceparent.py | 37 +++++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100644 src/mistralai/client/_hooks/traceparent.py diff --git a/src/mistralai/client/_hooks/registration.py b/src/mistralai/client/_hooks/registration.py index 58bebab0..de781de7 100644 --- a/src/mistralai/client/_hooks/registration.py +++ b/src/mistralai/client/_hooks/registration.py @@ -1,5 +1,6 @@ from .custom_user_agent import CustomUserAgentHook from .deprecation_warning import DeprecationWarningHook +from .traceparent import TraceparentInjectionHook from .tracing import TracingHook from .types import Hooks @@ -16,6 +17,7 @@ def init_hooks(hooks: Hooks): """ tracing_hook = TracingHook() hooks.register_before_request_hook(CustomUserAgentHook()) + hooks.register_before_request_hook(TraceparentInjectionHook()) hooks.register_after_success_hook(DeprecationWarningHook()) hooks.register_after_success_hook(tracing_hook) hooks.register_before_request_hook(tracing_hook) diff --git a/src/mistralai/client/_hooks/traceparent.py b/src/mistralai/client/_hooks/traceparent.py new file mode 100644 index 00000000..7118c41a --- /dev/null +++ b/src/mistralai/client/_hooks/traceparent.py @@ -0,0 +1,37 @@ +import random +from typing import Dict, Union + +import httpx +from opentelemetry.propagate import inject + +from .types import BeforeRequestContext, BeforeRequestHook + + +class TraceparentInjectionHook(BeforeRequestHook): + """Inject a sampled W3C traceparent header on workflow execute requests. + + Forwards the current OTEL span context if one is active and sampled, + otherwise generates a fresh sampled traceparent. This ensures worker traces + are always recorded regardless of the caller's sampling configuration. + """ + + def before_request( + self, hook_ctx: BeforeRequestContext, request: httpx.Request + ) -> Union[httpx.Request, Exception]: + if not request.url.path.endswith("/execute"): + return request + + # Don't overwrite an explicitly provided traceparent. + if "traceparent" in request.headers: + return request + + carrier: Dict[str, str] = {} + inject(carrier) + traceparent = carrier.get("traceparent", "") + if not traceparent.endswith("-01"): + trace_id = random.getrandbits(128) + span_id = random.getrandbits(64) + traceparent = f"00-{trace_id:032x}-{span_id:016x}-01" + + request.headers["traceparent"] = traceparent + return request From 5800ed1ce8af806d781254e9e292659e075de609 Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 09:18:17 +0200 Subject: [PATCH 2/6] Add tests for TraceparentInjectionHook Covers: no-op on non-execute paths, sampled header injection, explicit header preservation, OTEL context propagation, fallback for unsampled or absent spans, and uniqueness of generated IDs. --- .../extra/tests/test_traceparent_hook.py | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 src/mistralai/extra/tests/test_traceparent_hook.py diff --git a/src/mistralai/extra/tests/test_traceparent_hook.py b/src/mistralai/extra/tests/test_traceparent_hook.py new file mode 100644 index 00000000..87d36686 --- /dev/null +++ b/src/mistralai/extra/tests/test_traceparent_hook.py @@ -0,0 +1,118 @@ +"""Tests for TraceparentInjectionHook.""" + +import re +import unittest +from unittest.mock import MagicMock + +import httpx +from opentelemetry import trace +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.sdk.trace.sampling import ALWAYS_OFF + +from mistralai.client._hooks.traceparent import TraceparentInjectionHook +from mistralai.client._hooks.types import BeforeRequestContext, HookContext + + +TRACEPARENT_RE = re.compile(r"^00-[0-9a-f]{32}-[0-9a-f]{16}-01$") + + +def _make_hook_ctx(operation_id: str = "execute_workflow") -> BeforeRequestContext: + ctx = HookContext( + config=MagicMock(), + base_url="https://api.mistral.ai", + operation_id=operation_id, + oauth2_scopes=None, + security_source=None, + ) + return BeforeRequestContext(ctx) + + +def _make_request(path: str, traceparent: str | None = None) -> httpx.Request: + headers = {} + if traceparent is not None: + headers["traceparent"] = traceparent + return httpx.Request("POST", f"https://api.mistral.ai{path}", headers=headers) + + +class TestTraceparentInjectionHook(unittest.TestCase): + def setUp(self): + self.hook = TraceparentInjectionHook() + + # --- paths that must NOT be touched --- + + def test_non_execute_path_is_unchanged(self): + req = _make_request("/v1/workflows/my-wf/executions") + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertNotIn("traceparent", result.headers) + + def test_root_path_is_unchanged(self): + req = _make_request("/v1/workflows/my-wf") + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertNotIn("traceparent", result.headers) + + # --- /execute path: header injected --- + + def test_execute_path_gets_sampled_traceparent(self): + req = _make_request("/v1/workflows/my-wf/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertIn("traceparent", result.headers) + self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) + + def test_execute_registration_path_gets_sampled_traceparent(self): + req = _make_request("/v1/workflows/registrations/reg-id/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertIn("traceparent", result.headers) + self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) + + def test_explicit_traceparent_is_not_overwritten(self): + explicit = "00-aabbccddeeff00112233445566778899-0102030405060708-01" + req = _make_request("/v1/workflows/my-wf/execute", traceparent=explicit) + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertEqual(result.headers["traceparent"], explicit) + + # --- OTEL context propagation --- + + def test_propagates_sampled_active_span(self): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("parent") as span: + req = _make_request("/v1/workflows/my-wf/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + + injected = result.headers["traceparent"] + self.assertTrue(injected.endswith("-01")) + trace_id_hex = f"{span.get_span_context().trace_id:032x}" + self.assertIn(trace_id_hex, injected) + + def test_generates_fresh_traceparent_when_no_active_span(self): + req = _make_request("/v1/workflows/my-wf/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) + + def test_generates_fresh_traceparent_when_span_is_unsampled(self): + provider = TracerProvider(sampler=ALWAYS_OFF) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("unsampled-parent"): + req = _make_request("/v1/workflows/my-wf/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + + injected = result.headers["traceparent"] + self.assertTrue(injected.endswith("-01"), f"Expected sampled, got: {injected}") + + def test_generated_traceparents_are_unique(self): + ids = set() + for _ in range(50): + req = _make_request("/v1/workflows/my-wf/execute") + result = self.hook.before_request(_make_hook_ctx(), req) + ids.add(result.headers["traceparent"]) + self.assertEqual(len(ids), 50) + + +if __name__ == "__main__": + unittest.main() From a0d88ae3c51dd24fd96ffeda867bf237dc30e739 Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 09:40:36 +0200 Subject: [PATCH 3/6] Address review comments - Shorten TraceparentInjectionHook docstring to one line - Remove module docstring from test file - Drop low-ROI uniqueness test --- src/mistralai/client/_hooks/traceparent.py | 7 +------ src/mistralai/extra/tests/test_traceparent_hook.py | 11 ----------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/src/mistralai/client/_hooks/traceparent.py b/src/mistralai/client/_hooks/traceparent.py index 7118c41a..3804ec93 100644 --- a/src/mistralai/client/_hooks/traceparent.py +++ b/src/mistralai/client/_hooks/traceparent.py @@ -8,12 +8,7 @@ class TraceparentInjectionHook(BeforeRequestHook): - """Inject a sampled W3C traceparent header on workflow execute requests. - - Forwards the current OTEL span context if one is active and sampled, - otherwise generates a fresh sampled traceparent. This ensures worker traces - are always recorded regardless of the caller's sampling configuration. - """ + """Inject a sampled traceparent on /execute requests so worker traces are always recorded.""" def before_request( self, hook_ctx: BeforeRequestContext, request: httpx.Request diff --git a/src/mistralai/extra/tests/test_traceparent_hook.py b/src/mistralai/extra/tests/test_traceparent_hook.py index 87d36686..7ed18b81 100644 --- a/src/mistralai/extra/tests/test_traceparent_hook.py +++ b/src/mistralai/extra/tests/test_traceparent_hook.py @@ -1,5 +1,3 @@ -"""Tests for TraceparentInjectionHook.""" - import re import unittest from unittest.mock import MagicMock @@ -105,14 +103,5 @@ def test_generates_fresh_traceparent_when_span_is_unsampled(self): injected = result.headers["traceparent"] self.assertTrue(injected.endswith("-01"), f"Expected sampled, got: {injected}") - def test_generated_traceparents_are_unique(self): - ids = set() - for _ in range(50): - req = _make_request("/v1/workflows/my-wf/execute") - result = self.hook.before_request(_make_hook_ctx(), req) - ids.add(result.headers["traceparent"]) - self.assertEqual(len(ids), 50) - - if __name__ == "__main__": unittest.main() From 08763f7449bd642eeec0855ff9ac4f543af9aed6 Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 09:43:24 +0200 Subject: [PATCH 4/6] Fix lint errors in test file Remove unused opentelemetry.trace import (ruff F401) and add isinstance assertions so pyright can narrow Union[Request, Exception] before accessing .headers. --- src/mistralai/extra/tests/test_traceparent_hook.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/mistralai/extra/tests/test_traceparent_hook.py b/src/mistralai/extra/tests/test_traceparent_hook.py index 7ed18b81..3459be2f 100644 --- a/src/mistralai/extra/tests/test_traceparent_hook.py +++ b/src/mistralai/extra/tests/test_traceparent_hook.py @@ -3,7 +3,6 @@ from unittest.mock import MagicMock import httpx -from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -43,11 +42,13 @@ def setUp(self): def test_non_execute_path_is_unchanged(self): req = _make_request("/v1/workflows/my-wf/executions") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertNotIn("traceparent", result.headers) def test_root_path_is_unchanged(self): req = _make_request("/v1/workflows/my-wf") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertNotIn("traceparent", result.headers) # --- /execute path: header injected --- @@ -55,12 +56,14 @@ def test_root_path_is_unchanged(self): def test_execute_path_gets_sampled_traceparent(self): req = _make_request("/v1/workflows/my-wf/execute") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertIn("traceparent", result.headers) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) def test_execute_registration_path_gets_sampled_traceparent(self): req = _make_request("/v1/workflows/registrations/reg-id/execute") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertIn("traceparent", result.headers) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) @@ -68,6 +71,7 @@ def test_explicit_traceparent_is_not_overwritten(self): explicit = "00-aabbccddeeff00112233445566778899-0102030405060708-01" req = _make_request("/v1/workflows/my-wf/execute", traceparent=explicit) result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertEqual(result.headers["traceparent"], explicit) # --- OTEL context propagation --- @@ -82,6 +86,7 @@ def test_propagates_sampled_active_span(self): req = _make_request("/v1/workflows/my-wf/execute") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) injected = result.headers["traceparent"] self.assertTrue(injected.endswith("-01")) trace_id_hex = f"{span.get_span_context().trace_id:032x}" @@ -90,6 +95,7 @@ def test_propagates_sampled_active_span(self): def test_generates_fresh_traceparent_when_no_active_span(self): req = _make_request("/v1/workflows/my-wf/execute") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) def test_generates_fresh_traceparent_when_span_is_unsampled(self): @@ -100,6 +106,7 @@ def test_generates_fresh_traceparent_when_span_is_unsampled(self): req = _make_request("/v1/workflows/my-wf/execute") result = self.hook.before_request(_make_hook_ctx(), req) + assert isinstance(result, httpx.Request) injected = result.headers["traceparent"] self.assertTrue(injected.endswith("-01"), f"Expected sampled, got: {injected}") From 0b6c2a79afa0ac621f85f1a153fc6f167dc795a1 Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 09:58:59 +0200 Subject: [PATCH 5/6] Use operation ID instead of URL path to identify execute calls Matching on request.url.path.endswith("/execute") would affect any future endpoint that happens to share that suffix. Keying on the operation ID is explicit and safe. --- src/mistralai/client/_hooks/traceparent.py | 8 +++- .../extra/tests/test_traceparent_hook.py | 38 +++++++++---------- 2 files changed, 25 insertions(+), 21 deletions(-) diff --git a/src/mistralai/client/_hooks/traceparent.py b/src/mistralai/client/_hooks/traceparent.py index 3804ec93..a8ddedf4 100644 --- a/src/mistralai/client/_hooks/traceparent.py +++ b/src/mistralai/client/_hooks/traceparent.py @@ -7,13 +7,19 @@ from .types import BeforeRequestContext, BeforeRequestHook +_EXECUTE_OPERATION_IDS = { + "execute_workflow_v1_workflows__workflow_identifier__execute_post", + "execute_workflow_registration_v1_workflows_registrations__workflow_registration_id__execute_post", +} + + class TraceparentInjectionHook(BeforeRequestHook): """Inject a sampled traceparent on /execute requests so worker traces are always recorded.""" def before_request( self, hook_ctx: BeforeRequestContext, request: httpx.Request ) -> Union[httpx.Request, Exception]: - if not request.url.path.endswith("/execute"): + if hook_ctx.operation_id not in _EXECUTE_OPERATION_IDS: return request # Don't overwrite an explicitly provided traceparent. diff --git a/src/mistralai/extra/tests/test_traceparent_hook.py b/src/mistralai/extra/tests/test_traceparent_hook.py index 3459be2f..9341b6c7 100644 --- a/src/mistralai/extra/tests/test_traceparent_hook.py +++ b/src/mistralai/extra/tests/test_traceparent_hook.py @@ -8,14 +8,18 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.trace.sampling import ALWAYS_OFF -from mistralai.client._hooks.traceparent import TraceparentInjectionHook +from mistralai.client._hooks.traceparent import TraceparentInjectionHook, _EXECUTE_OPERATION_IDS from mistralai.client._hooks.types import BeforeRequestContext, HookContext TRACEPARENT_RE = re.compile(r"^00-[0-9a-f]{32}-[0-9a-f]{16}-01$") +_EXECUTE_OP_ID = "execute_workflow_v1_workflows__workflow_identifier__execute_post" +_EXECUTE_REG_OP_ID = "execute_workflow_registration_v1_workflows_registrations__workflow_registration_id__execute_post" +_OTHER_OP_ID = "list_executions_v1_workflows__workflow_identifier__executions_get" -def _make_hook_ctx(operation_id: str = "execute_workflow") -> BeforeRequestContext: + +def _make_hook_ctx(operation_id: str = _EXECUTE_OP_ID) -> BeforeRequestContext: ctx = HookContext( config=MagicMock(), base_url="https://api.mistral.ai", @@ -37,32 +41,26 @@ class TestTraceparentInjectionHook(unittest.TestCase): def setUp(self): self.hook = TraceparentInjectionHook() - # --- paths that must NOT be touched --- + # --- non-execute operations must NOT be touched --- - def test_non_execute_path_is_unchanged(self): + def test_other_operation_is_unchanged(self): req = _make_request("/v1/workflows/my-wf/executions") - result = self.hook.before_request(_make_hook_ctx(), req) - assert isinstance(result, httpx.Request) - self.assertNotIn("traceparent", result.headers) - - def test_root_path_is_unchanged(self): - req = _make_request("/v1/workflows/my-wf") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_OTHER_OP_ID), req) assert isinstance(result, httpx.Request) self.assertNotIn("traceparent", result.headers) - # --- /execute path: header injected --- + # --- execute operations: header injected --- - def test_execute_path_gets_sampled_traceparent(self): + def test_execute_gets_sampled_traceparent(self): req = _make_request("/v1/workflows/my-wf/execute") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_OP_ID), req) assert isinstance(result, httpx.Request) self.assertIn("traceparent", result.headers) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) - def test_execute_registration_path_gets_sampled_traceparent(self): + def test_execute_registration_gets_sampled_traceparent(self): req = _make_request("/v1/workflows/registrations/reg-id/execute") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_REG_OP_ID), req) assert isinstance(result, httpx.Request) self.assertIn("traceparent", result.headers) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) @@ -70,7 +68,7 @@ def test_execute_registration_path_gets_sampled_traceparent(self): def test_explicit_traceparent_is_not_overwritten(self): explicit = "00-aabbccddeeff00112233445566778899-0102030405060708-01" req = _make_request("/v1/workflows/my-wf/execute", traceparent=explicit) - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_OP_ID), req) assert isinstance(result, httpx.Request) self.assertEqual(result.headers["traceparent"], explicit) @@ -84,7 +82,7 @@ def test_propagates_sampled_active_span(self): with tracer.start_as_current_span("parent") as span: req = _make_request("/v1/workflows/my-wf/execute") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_OP_ID), req) assert isinstance(result, httpx.Request) injected = result.headers["traceparent"] @@ -94,7 +92,7 @@ def test_propagates_sampled_active_span(self): def test_generates_fresh_traceparent_when_no_active_span(self): req = _make_request("/v1/workflows/my-wf/execute") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_OP_ID), req) assert isinstance(result, httpx.Request) self.assertRegex(result.headers["traceparent"], TRACEPARENT_RE) @@ -104,7 +102,7 @@ def test_generates_fresh_traceparent_when_span_is_unsampled(self): with tracer.start_as_current_span("unsampled-parent"): req = _make_request("/v1/workflows/my-wf/execute") - result = self.hook.before_request(_make_hook_ctx(), req) + result = self.hook.before_request(_make_hook_ctx(_EXECUTE_OP_ID), req) assert isinstance(result, httpx.Request) injected = result.headers["traceparent"] From 6bc4503e3a7dee4f21a2b0d882ff2a8c432b487f Mon Sep 17 00:00:00 2001 From: arnaudde Date: Tue, 31 Mar 2026 10:04:44 +0200 Subject: [PATCH 6/6] Remove unused _EXECUTE_OPERATION_IDS import --- src/mistralai/extra/tests/test_traceparent_hook.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mistralai/extra/tests/test_traceparent_hook.py b/src/mistralai/extra/tests/test_traceparent_hook.py index 9341b6c7..8202f3d6 100644 --- a/src/mistralai/extra/tests/test_traceparent_hook.py +++ b/src/mistralai/extra/tests/test_traceparent_hook.py @@ -8,7 +8,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.trace.sampling import ALWAYS_OFF -from mistralai.client._hooks.traceparent import TraceparentInjectionHook, _EXECUTE_OPERATION_IDS +from mistralai.client._hooks.traceparent import TraceparentInjectionHook from mistralai.client._hooks.types import BeforeRequestContext, HookContext