From f1f57ba4c560ffe2b572af36fccd34cbcec6d82e Mon Sep 17 00:00:00 2001
From: James Hargreaves <james.hargreaves@portswigger.net>
Date: Mon, 22 Jun 2026 15:37:00 +0100
Subject: [PATCH] test(posthog): add coverage for cache token fields in PostHog
 callback and _inject_cache_tokens dict path

---
 tests/logging_callback_tests/test_posthog.py  | 37 +++++++++++++++++++
 .../test_standard_logging_payload.py          | 16 ++++++++
 2 files changed, 53 insertions(+)

diff --git a/tests/logging_callback_tests/test_posthog.py b/tests/logging_callback_tests/test_posthog.py
index b3f346bcf9d..c32aa6c9b05 100644
--- a/tests/logging_callback_tests/test_posthog.py
+++ b/tests/logging_callback_tests/test_posthog.py
@@ -569,3 +569,40 @@ async def test_sync_callback_not_affected_by_atexit():
             assert (
                 callback_invoked_immediately
             ), "Sync callback should be invoked immediately"
+
+
+@pytest.mark.asyncio
+async def test_cache_tokens_emitted_when_nonzero():
+    """$ai_cache_read_input_tokens / $ai_cache_creation_input_tokens are
+    conditionally emitted — only when the corresponding field is non-zero."""
+    from create_mock_standard_logging_payload import create_standard_logging_payload
+
+    posthog_logger = PostHogLogger()
+    standard_payload = create_standard_logging_payload()
+    standard_payload["cache_read_input_tokens"] = 500
+    standard_payload["cache_creation_input_tokens"] = 100
+    kwargs = {"standard_logging_object": standard_payload}
+
+    event_payload = posthog_logger.create_posthog_event_payload(kwargs)
+    props = event_payload["properties"]
+
+    assert props["$ai_cache_read_input_tokens"] == 500
+    assert props["$ai_cache_creation_input_tokens"] == 100
+
+
+@pytest.mark.asyncio
+async def test_cache_tokens_absent_when_zero():
+    """Cache token fields must not be emitted when the value is 0 —
+    keeps PostHog events clean for non-caching providers."""
+    from create_mock_standard_logging_payload import create_standard_logging_payload
+
+    posthog_logger = PostHogLogger()
+    standard_payload = create_standard_logging_payload()
+    # Both default to 0; confirm they are not forwarded to PostHog.
+    kwargs = {"standard_logging_object": standard_payload}
+
+    event_payload = posthog_logger.create_posthog_event_payload(kwargs)
+    props = event_payload["properties"]
+
+    assert "$ai_cache_read_input_tokens" not in props
+    assert "$ai_cache_creation_input_tokens" not in props
diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py
index 34098d63a90..1c1d807e825 100644
--- a/tests/logging_callback_tests/test_standard_logging_payload.py
+++ b/tests/logging_callback_tests/test_standard_logging_payload.py
@@ -1184,3 +1184,19 @@ def test_inject_cache_tokens_prefers_private_attr_then_prompt_tokens_details():
     d = StandardLoggingPayloadSetup._inject_cache_tokens({}, openai)
     assert d["cache_read_input_tokens"] == 80
     assert d["cache_creation_input_tokens"] == 0
+
+
+def test_inject_cache_tokens_prompt_tokens_details_as_dict():
+    """_inject_cache_tokens must handle prompt_tokens_details as a plain dict
+    (the branch distinct from the PromptTokensDetailsWrapper object path)."""
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+    from litellm.types.utils import Usage
+
+    # Construct a Usage where prompt_tokens_details is stored as a dict.
+    usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+    # Directly set a dict to exercise the isinstance(ptd, dict) branch.
+    object.__setattr__(usage, "prompt_tokens_details", {"cached_tokens": 42})
+
+    d = StandardLoggingPayloadSetup._inject_cache_tokens({}, usage)
+    assert d["cache_read_input_tokens"] == 42
+    assert d["cache_creation_input_tokens"] == 0