fix(claude-messages): put the input content write inside the guard too

The prompt write ran before the try that fails the span it writes to. Serialising
conversation content raises on anything that is not JSON-serialisable, so a raise
there left the chat span open in the tool loop, and on both root paths left the
root open: never ended, never exported, so the run disappeared from AI Config
Monitoring along with the feature_flag event it carries.

The output writes were moved inside their guards earlier in this stack. The input
writes were not, which is the same defect at the other end of the same span.

Two tests, one per root path.

Found by Bugbot on #34, which is this shape in langchain-messages.

Author: apucacao

packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py

@@ -137,16 +137,6 @@ async def _run_tool_loop(
 
     while True:
         model_span = start_model_span(config, parent)
-        # Written before the call, so an in-flight or failed turn still shows what it was asked.
-        # `conversation` grows with each turn, which is what makes a `chat` span self-contained.
-        if capture_content:
-            set_input_content_attributes(
-                model_span,
-                capture_content,
-                system_instructions=system,
-                messages=to_span_messages(conversation),
-                tool_definitions=tool_definitions,
-            )
 
         kwargs: dict[str, Any] = {
             "model": config["model"]["name"],
@@ -163,6 +153,18 @@ async def _run_tool_loop(
         # span open: the blocking path has no `finally` that could recover it, so the turn would never
         # be exported and the run would show a model call that left no trace.
         try:
+            # Written before the call, so an in-flight or failed turn still shows what it was asked.
+            # `conversation` grows with each turn, which is what makes a `chat` span self-contained.
+            # Inside the guard, because serialising it raises on anything that is not
+            # JSON-serialisable and a raise out here would leave this span open.
+            if capture_content:
+                set_input_content_attributes(
+                    model_span,
+                    capture_content,
+                    system_instructions=system,
+                    messages=to_span_messages(conversation),
+                    tool_definitions=tool_definitions,
+                )
             resp = await client.messages.create(**kwargs)
 
             raw_usage = raw_usage_of(getattr(resp, "usage", None))
@@ -274,17 +276,20 @@ async def _call_impl(
         parent = parent_context_of(span)
 
         messages, system = _build_messages(config, user_input, vs, history=history)
-        set_input_content_attributes(
-            span,
-            capture_content,
-            system_instructions=system,
-            messages=to_span_messages(messages),
-        )
 
         # Outside the try, so the failure path can still report the spend of the turns that
         # completed before it.
         run_usage = RawRunUsage()
         try:
+            # Inside the guard, because serialising the prompt raises on anything that is not
+            # JSON-serialisable and a raise out here would leave the root open: never ended, never
+            # exported, and the run gone from AI Config Monitoring with the feature_flag event on it.
+            set_input_content_attributes(
+                span,
+                capture_content,
+                system_instructions=system,
+                messages=to_span_messages(messages),
+            )
             output, usage = await _run_tool_loop(
                 client,
                 config,
@@ -364,12 +369,6 @@ async def _stream_gen(
     messages, system = _build_messages(
         config, user_input, variables, include_output_format=False, history=history
     )
-    set_input_content_attributes(
-        span,
-        capture_content,
-        system_instructions=system,
-        messages=to_span_messages(messages),
-    )
 
     tools = _build_tools(config.get("tools") or {})
     tool_definitions = to_tool_definitions(tools)
@@ -389,6 +388,15 @@ async def _stream_gen(
     run_usage = RawRunUsage()
 
     try:
+        # Inside the guard, because serialising the prompt raises on anything that is not
+        # JSON-serialisable. A raise out here would leave the root open with the `finally` never
+        # entered, so the run would vanish from AI Config Monitoring with its feature_flag event.
+        set_input_content_attributes(
+            span,
+            capture_content,
+            system_instructions=system,
+            messages=to_span_messages(messages),
+        )
         while True:
             model_span = start_model_span(config, parent)
             open_model_span = model_span

packages/claude-messages/tests/test_handler.py

@@ -1968,3 +1968,56 @@ async def _ctx_mgr() -> AsyncGenerator[Any, None]:
 
         assert rec.root.attributes["gen_ai.usage.input_tokens"] == 37
         assert rec.root.attributes["gen_ai.usage.output_tokens"] == 11
+
+
+class TestInputWritesNeverLeakASpan:
+    """Serialising the prompt must not be able to strand a span.
+
+    The input content write ran before the guard that fails the span it writes to, so a raise there
+    left the root open on both paths: never ended, never exported, so the run disappeared from AI
+    Config Monitoring along with the feature_flag event it carries.
+    """
+
+    async def test_the_blocking_root_still_ends(
+        self, mock_anthropic: MagicMock
+    ) -> None:
+        from opentelemetry.trace import StatusCode
+
+        from launchdarkly_ai_claude_messages import create_claude_messages_handler
+
+        ctx, rec = _recording()
+        with (
+            ctx,
+            patch(
+                "launchdarkly_ai_claude_messages.handler.set_input_content_attributes",
+                side_effect=TypeError("cannot serialise this prompt"),
+            ),
+            pytest.raises(TypeError),
+        ):
+            await create_claude_messages_handler(capture_content=True)(
+                CONFIG, "q", {}, {}
+            )
+
+        assert rec.root.ended == 1, "the root span leaked"
+        assert StatusCode.ERROR in rec.root.statuses
+
+    async def test_the_streaming_root_still_ends(
+        self, mock_anthropic: MagicMock
+    ) -> None:
+        from launchdarkly_ai_claude_messages import create_claude_messages_handler
+
+        ctx, rec = _recording()
+        with (
+            ctx,
+            patch(
+                "launchdarkly_ai_claude_messages.handler.set_input_content_attributes",
+                side_effect=TypeError("cannot serialise this prompt"),
+            ),
+            pytest.raises(TypeError),
+        ):
+            async for _ in await create_claude_messages_handler(
+                capture_content=True
+            ).stream(CONFIG, "q"):
+                pass
+
+        assert rec.root.ended == 1, "the root span leaked"