fix(llma): two $ai_generation fidelity bugs in the Claude Code session ingest

posthog_llma/event_builder.py

@@ -167,7 +167,13 @@ def build_events(parsed: dict, config: dict) -> list[dict]:
         if not privacy_mode:
             max_attr = config.get("max_attribute_length", 12000)
             content_blocks = []
-            if gen["output_text"]:
+            # Prefer the structured blocks so extended-thinking content stays
+            # typed as "thinking" rather than being flattened into assistant
+            # text. Older parsed payloads have no output_blocks, so fall back
+            # to the flattened output_text.
+            if gen.get("output_blocks"):
+                content_blocks.extend(gen["output_blocks"])
+            elif gen["output_text"]:
                 content_blocks.append({"type": "text", "text": gen["output_text"]})
             content_blocks.extend(_truncate_tool_blocks(gen.get("tool_use_blocks", []), max_attr))
             if content_blocks:

posthog_llma/events.py

@@ -69,8 +69,8 @@ def build_ai_generation(
         "$ai_framework": "claude-code",
         "$ai_project_name": project_name,
         "$ai_agent_name": agent_name,
-        "cache_read_input_tokens": cache_read_tokens,
-        "cache_creation_input_tokens": cache_creation_tokens,
+        "$ai_cache_read_input_tokens": cache_read_tokens,
+        "$ai_cache_creation_input_tokens": cache_creation_tokens,
     }
 
     if user_prompt and not privacy_mode:

posthog_llma/parser.py

@@ -332,18 +332,24 @@ def _finalize_generation(state: dict) -> tuple[dict, list]:
     type_order = state.get("type_order") or ["thinking", "text", "tool_use"]
 
     text_parts = []
+    # Structured counterpart to text_parts: keeps thinking and text distinct so
+    # downstream consumers can render/filter them separately. `output_text`
+    # remains the flattened join for backward compatibility.
+    output_blocks = []
     entry_tool_uses = []
     for block_type in type_order:
         if block_type == "thinking":
             for item in blocks.get("thinking", []):
                 t = item.get("thinking", "")
                 if t:
                     text_parts.append(t)
+                    output_blocks.append({"type": "thinking", "thinking": t})
         elif block_type == "text":
             for item in blocks.get("text", []):
                 t = item.get("text", "")
                 if t:
                     text_parts.append(t)
+                    output_blocks.append({"type": "text", "text": t})
         elif block_type == "tool_use":
             for item in blocks.get("tool_use", []):
                 entry_tool_uses.append({
@@ -374,6 +380,7 @@ def _finalize_generation(state: dict) -> tuple[dict, list]:
         "span_id": span_id,
         "msg_id": state.get("msg_id", ""),
         "output_text": output_text,
+        "output_blocks": output_blocks,
         "tool_use_blocks": tool_use_blocks,
         "is_error": stop_reason == "error",
         "error_message": state.get("error_message"),

tests/test_posthog_llma.py

@@ -83,6 +83,43 @@ def test_timestamp_passthrough(self):
         )
         assert ev["timestamp"] == "2026-04-12T21:00:00Z"
 
+    def test_cache_tokens_use_ai_namespace(self):
+        """Cache token properties must carry the $ai_ prefix.
+
+        PostHog's LLM Analytics pipeline reads $ai_cache_read_input_tokens and
+        $ai_cache_creation_input_tokens when computing cost. Unprefixed keys are
+        ingested as ordinary custom properties and ignored by cost calculation,
+        which silently understates spend on prompt-cached workloads.
+        """
+        ev = build_ai_generation(
+            model="claude-opus-4-6",
+            input_tokens=4,
+            output_tokens=926,
+            cache_read_tokens=150109,
+            cache_creation_tokens=75729,
+            trace_id="t", session_id="s",
+        )
+        props = ev["properties"]
+        assert props["$ai_cache_read_input_tokens"] == 150109
+        assert props["$ai_cache_creation_input_tokens"] == 75729
+        # The unprefixed keys must not linger, or the same numbers get ingested
+        # twice under two different names.
+        assert "cache_read_input_tokens" not in props
+        assert "cache_creation_input_tokens" not in props
+
+    def test_cache_tokens_emitted_even_when_zero(self):
+        """Anthropic generations always carry cache fields, matching posthog-python.
+
+        See posthog/ai/utils.py: for the anthropic provider the official SDK
+        always includes both cache fields even at 0, rather than omitting them.
+        """
+        ev = build_ai_generation(
+            model="m", trace_id="t", session_id="s",
+        )
+        props = ev["properties"]
+        assert props["$ai_cache_read_input_tokens"] == 0
+        assert props["$ai_cache_creation_input_tokens"] == 0
+
     def test_no_timestamp_means_no_key(self):
         ev = build_ai_generation(
             model="m", trace_id="t", session_id="s",
@@ -96,8 +133,8 @@ def test_cache_tokens(self):
             cache_read_tokens=5, cache_creation_tokens=3,
         )
         props = ev["properties"]
-        assert props["cache_read_input_tokens"] == 5
-        assert props["cache_creation_input_tokens"] == 3
+        assert props["$ai_cache_read_input_tokens"] == 5
+        assert props["$ai_cache_creation_input_tokens"] == 3
 
     def test_no_cost_properties(self):
         """Cost is calculated by PostHog ingestion, we should not send it."""

tests/test_session_parser.py

@@ -314,12 +314,14 @@ def test_thinking_preserved_when_split_across_chunks(self):
             assert len(parsed["tool_uses"]) == 1
             assert parsed["tool_uses"][0]["name"] == "Bash"
 
-            # Downstream events include thinking in output_choices
+            # Downstream events include thinking in output_choices, typed as a
+            # thinking block rather than flattened into assistant text.
             events = build_events(parsed, DEFAULT_CONFIG)
             gen_ev = next(e for e in events if e["event"] == "$ai_generation")
             content_blocks = gen_ev["properties"]["$ai_output_choices"][0]["content"]
-            text_blocks = [b for b in content_blocks if b.get("type") == "text"]
-            assert text_blocks and "deep thoughts" in text_blocks[0]["text"]
+            thinking_blocks = [b for b in content_blocks if b.get("type") == "thinking"]
+            assert thinking_blocks and thinking_blocks[0]["thinking"] == "deep thoughts"
+            assert not [b for b in content_blocks if b.get("type") == "text"]
         finally:
             os.unlink(path)
 
@@ -514,6 +516,65 @@ def test_merge_preserves_first_seen_block_order(self):
             # Text came first in the source, so output_text must lead
             # with "preface" rather than reshuffling thinking ahead of it.
             assert parsed["generations"][0]["output_text"] == "preface\nafterthought"
+            # The same order is preserved in the structured blocks, with each
+            # block keeping its own type rather than collapsing into text.
+            blocks = parsed["generations"][0]["output_blocks"]
+            assert blocks == [
+                {"type": "text", "text": "preface"},
+                {"type": "thinking", "thinking": "afterthought"},
+            ]
+        finally:
+            os.unlink(path)
+
+    def test_thinking_and_text_stay_separate_blocks(self):
+        """Extended thinking must not be ingested as assistant text.
+
+        _finalize_generation joins thinking and text into a single output_text
+        for backward compatibility, but $ai_output_choices needs them typed
+        separately or thinking content renders as ordinary assistant output and
+        cannot be filtered on.
+        """
+        entries = [
+            {"type": "permission-mode", "permissionMode": "default", "sessionId": "s1"},
+            {
+                "type": "user", "uuid": "u-1", "parentUuid": None,
+                "promptId": "p-1", "isMeta": False,
+                "message": {"role": "user", "content": "go"},
+                "timestamp": "2026-04-12T10:00:00.000Z",
+                "sessionId": "s1", "version": "2.1.0", "cwd": "/tmp",
+            },
+            {
+                "type": "assistant", "uuid": "a-1", "parentUuid": "u-1",
+                "message": {
+                    "role": "assistant", "id": "msg-1",
+                    "model": "claude-opus-4-6", "stop_reason": "end_turn",
+                    "usage": {"input_tokens": 5, "output_tokens": 10},
+                    "content": [
+                        {"type": "thinking", "thinking": "weighing the options"},
+                        {"type": "text", "text": "Here is the answer."},
+                    ],
+                },
+                "timestamp": "2026-04-12T10:00:01.000Z",
+                "sessionId": "s1", "version": "2.1.0", "cwd": "/tmp",
+            },
+        ]
+        path = _write_jsonl(entries)
+        try:
+            parsed = parse_session(path, DEFAULT_CONFIG)
+            gen = parsed["generations"][0]
+            # Flattened form is unchanged.
+            assert gen["output_text"] == "weighing the options\nHere is the answer."
+
+            events = build_events(parsed, DEFAULT_CONFIG)
+            gen_ev = next(e for e in events if e["event"] == "$ai_generation")
+            content = gen_ev["properties"]["$ai_output_choices"][0]["content"]
+
+            thinking = [b for b in content if b.get("type") == "thinking"]
+            text = [b for b in content if b.get("type") == "text"]
+            assert len(thinking) == 1 and thinking[0]["thinking"] == "weighing the options"
+            assert len(text) == 1 and text[0]["text"] == "Here is the answer."
+            # And the reasoning must not leak into the visible answer text.
+            assert "weighing the options" not in text[0]["text"]
         finally:
             os.unlink(path)