posthoganalytics 7.6.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

posthoganalytics/test/ai/openai_agents/test_processor.py

@@ -0,0 +1,810 @@
+import logging
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+try:
+    from agents.tracing.span_data import (
+        AgentSpanData,
+        CustomSpanData,
+        FunctionSpanData,
+        GenerationSpanData,
+        GuardrailSpanData,
+        HandoffSpanData,
+        ResponseSpanData,
+        SpeechSpanData,
+        TranscriptionSpanData,
+    )
+
+    from posthoganalytics.ai.openai_agents import PostHogTracingProcessor, instrument
+
+    OPENAI_AGENTS_AVAILABLE = True
+except ImportError:
+    OPENAI_AGENTS_AVAILABLE = False
+
+
+# Skip all tests if OpenAI Agents SDK is not available
+pytestmark = pytest.mark.skipif(
+    not OPENAI_AGENTS_AVAILABLE, reason="OpenAI Agents SDK is not available"
+)
+
+
+@pytest.fixture(scope="function")
+def mock_client():
+    client = MagicMock()
+    client.privacy_mode = False
+    logging.getLogger("posthog").setLevel(logging.DEBUG)
+    return client
+
+
+@pytest.fixture(scope="function")
+def processor(mock_client):
+    return PostHogTracingProcessor(
+        client=mock_client,
+        distinct_id="test-user",
+        privacy_mode=False,
+    )
+
+
+@pytest.fixture
+def mock_trace():
+    trace = MagicMock()
+    trace.trace_id = "trace_123456789"
+    trace.name = "Test Workflow"
+    trace.group_id = "group_123"
+    trace.metadata = {"key": "value"}
+    return trace
+
+
+@pytest.fixture
+def mock_span():
+    span = MagicMock()
+    span.trace_id = "trace_123456789"
+    span.span_id = "span_987654321"
+    span.parent_id = None
+    span.started_at = "2024-01-01T00:00:00Z"
+    span.ended_at = "2024-01-01T00:00:01Z"
+    span.error = None
+    return span
+
+
+class TestPostHogTracingProcessor:
+    """Tests for the PostHogTracingProcessor class."""
+
+    def test_initialization(self, mock_client):
+        """Test processor initializes correctly."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="user@example.com",
+            privacy_mode=True,
+            groups={"company": "acme"},
+            properties={"env": "test"},
+        )
+
+        assert processor._client == mock_client
+        assert processor._distinct_id == "user@example.com"
+        assert processor._privacy_mode is True
+        assert processor._groups == {"company": "acme"}
+        assert processor._properties == {"env": "test"}
+
+    def test_initialization_with_callable_distinct_id(self, mock_client, mock_trace):
+        """Test processor with callable distinct_id resolver."""
+
+        def resolver(trace):
+            return trace.metadata.get("user_id", "default")
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        mock_trace.metadata = {"user_id": "resolved-user"}
+        distinct_id = processor._get_distinct_id(mock_trace)
+        assert distinct_id == "resolved-user"
+
+    def test_on_trace_start_stores_metadata(self, processor, mock_client, mock_trace):
+        """Test that on_trace_start stores metadata but does not capture an event."""
+        processor.on_trace_start(mock_trace)
+
+        mock_client.capture.assert_not_called()
+        assert mock_trace.trace_id in processor._trace_metadata
+
+    def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace):
+        """Test that on_trace_end captures $ai_trace event."""
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        mock_client.capture.assert_called_once()
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_trace"
+        assert call_kwargs["distinct_id"] == "test-user"
+        assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
+        assert call_kwargs["properties"]["$ai_trace_name"] == "Test Workflow"
+        assert call_kwargs["properties"]["$ai_provider"] == "openai"
+        assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
+        assert "$ai_latency" in call_kwargs["properties"]
+
+    def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace):
+        """Test that trace events use personless mode when no distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        # Should fallback to trace_id as the distinct_id
+        assert call_kwargs["distinct_id"] == mock_trace.trace_id
+
+    def test_personless_mode_for_spans_when_no_distinct_id(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test that span events use personless mode when no distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+        )
+
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        assert call_kwargs["distinct_id"] == mock_span.trace_id
+
+    def test_personless_mode_when_callable_returns_none(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test personless mode when callable distinct_id returns None."""
+
+        def resolver(trace):
+            return None  # Simulate no user ID available
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        assert call_kwargs["distinct_id"] == mock_span.trace_id
+
+    def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace):
+        """Test that events create person profiles when distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="real-user",
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert "$process_person_profile" not in call_kwargs["properties"]
+
+    def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace):
+        """Test that on_trace_end clears stored trace metadata."""
+        processor.on_trace_start(mock_trace)
+        assert mock_trace.trace_id in processor._trace_metadata
+
+        processor.on_trace_end(mock_trace)
+        assert mock_trace.trace_id not in processor._trace_metadata
+        # Also verify it captured the event
+        mock_client.capture.assert_called_once()
+
+    def test_on_span_start_tracks_time(self, processor, mock_span):
+        """Test that on_span_start records start time."""
+        processor.on_span_start(mock_span)
+        assert mock_span.span_id in processor._span_start_times
+
+    def test_generation_span_mapping(self, processor, mock_client, mock_span):
+        """Test GenerationSpanData maps to $ai_generation event."""
+        span_data = GenerationSpanData(
+            input=[{"role": "user", "content": "Hello"}],
+            output=[{"role": "assistant", "content": "Hi there!"}],
+            model="gpt-4o",
+            model_config={"temperature": 0.7, "max_tokens": 100},
+            usage={"input_tokens": 10, "output_tokens": 20},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        mock_client.capture.assert_called_once()
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_generation"
+        assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
+        assert call_kwargs["properties"]["$ai_span_id"] == "span_987654321"
+        assert call_kwargs["properties"]["$ai_provider"] == "openai"
+        assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
+        assert call_kwargs["properties"]["$ai_model"] == "gpt-4o"
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 10
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 20
+        assert call_kwargs["properties"]["$ai_input"] == [
+            {"role": "user", "content": "Hello"}
+        ]
+        assert call_kwargs["properties"]["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Hi there!"}
+        ]
+
+    def test_generation_span_with_reasoning_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test GenerationSpanData includes reasoning tokens when present."""
+        span_data = GenerationSpanData(
+            model="o1-preview",
+            usage={
+                "input_tokens": 100,
+                "output_tokens": 500,
+                "reasoning_tokens": 400,
+            },
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400
+
+    def test_function_span_mapping(self, processor, mock_client, mock_span):
+        """Test FunctionSpanData maps to $ai_span event with type=tool."""
+        span_data = FunctionSpanData(
+            name="get_weather",
+            input='{"city": "San Francisco"}',
+            output="Sunny, 72F",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "get_weather"
+        assert call_kwargs["properties"]["$ai_span_type"] == "tool"
+        assert (
+            call_kwargs["properties"]["$ai_input_state"] == '{"city": "San Francisco"}'
+        )
+        assert call_kwargs["properties"]["$ai_output_state"] == "Sunny, 72F"
+
+    def test_agent_span_mapping(self, processor, mock_client, mock_span):
+        """Test AgentSpanData maps to $ai_span event with type=agent."""
+        span_data = AgentSpanData(
+            name="CustomerServiceAgent",
+            handoffs=["TechnicalAgent", "BillingAgent"],
+            tools=["search", "get_order"],
+            output_type="str",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "CustomerServiceAgent"
+        assert call_kwargs["properties"]["$ai_span_type"] == "agent"
+        assert call_kwargs["properties"]["$ai_agent_handoffs"] == [
+            "TechnicalAgent",
+            "BillingAgent",
+        ]
+        assert call_kwargs["properties"]["$ai_agent_tools"] == ["search", "get_order"]
+
+    def test_handoff_span_mapping(self, processor, mock_client, mock_span):
+        """Test HandoffSpanData maps to $ai_span event with type=handoff."""
+        span_data = HandoffSpanData(
+            from_agent="TriageAgent",
+            to_agent="TechnicalAgent",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "handoff"
+        assert call_kwargs["properties"]["$ai_handoff_from_agent"] == "TriageAgent"
+        assert call_kwargs["properties"]["$ai_handoff_to_agent"] == "TechnicalAgent"
+        assert (
+            call_kwargs["properties"]["$ai_span_name"]
+            == "TriageAgent -> TechnicalAgent"
+        )
+
+    def test_guardrail_span_mapping(self, processor, mock_client, mock_span):
+        """Test GuardrailSpanData maps to $ai_span event with type=guardrail."""
+        span_data = GuardrailSpanData(
+            name="ContentFilter",
+            triggered=True,
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "ContentFilter"
+        assert call_kwargs["properties"]["$ai_span_type"] == "guardrail"
+        assert call_kwargs["properties"]["$ai_guardrail_triggered"] is True
+
+    def test_custom_span_mapping(self, processor, mock_client, mock_span):
+        """Test CustomSpanData maps to $ai_span event with type=custom."""
+        span_data = CustomSpanData(
+            name="database_query",
+            data={"query": "SELECT * FROM users", "rows": 100},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "database_query"
+        assert call_kwargs["properties"]["$ai_span_type"] == "custom"
+        assert call_kwargs["properties"]["$ai_custom_data"] == {
+            "query": "SELECT * FROM users",
+            "rows": 100,
+        }
+
+    def test_privacy_mode_redacts_content(self, mock_client, mock_span):
+        """Test that privacy_mode redacts input/output content."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            privacy_mode=True,
+        )
+
+        span_data = GenerationSpanData(
+            input=[{"role": "user", "content": "Secret message"}],
+            output=[{"role": "assistant", "content": "Secret response"}],
+            model="gpt-4o",
+            usage={"input_tokens": 10, "output_tokens": 20},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        # Content should be redacted
+        assert call_kwargs["properties"]["$ai_input"] is None
+        assert call_kwargs["properties"]["$ai_output_choices"] is None
+        # Token counts should still be present
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 10
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 20
+
+    def test_error_handling_in_span(self, processor, mock_client, mock_span):
+        """Test that span errors are captured correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "Rate limit exceeded", "data": {"code": 429}}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["properties"]["$ai_is_error"] is True
+        assert call_kwargs["properties"]["$ai_error"] == "Rate limit exceeded"
+
+    def test_generation_span_includes_total_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that $ai_total_tokens is calculated and included."""
+        span_data = GenerationSpanData(
+            model="gpt-4o",
+            usage={"input_tokens": 100, "output_tokens": 50},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 150
+
+    def test_error_type_categorization_model_behavior(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that ModelBehaviorError is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "ModelBehaviorError: Invalid JSON output",
+            "type": "ModelBehaviorError",
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
+
+    def test_error_type_categorization_user_error(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that UserError is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "UserError: Tool failed", "type": "UserError"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "user_error"
+
+    def test_error_type_categorization_input_guardrail(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that InputGuardrailTripwireTriggered is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "InputGuardrailTripwireTriggered: Content blocked"
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert (
+            call_kwargs["properties"]["$ai_error_type"] == "input_guardrail_triggered"
+        )
+
+    def test_error_type_categorization_output_guardrail(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that OutputGuardrailTripwireTriggered is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "OutputGuardrailTripwireTriggered: Response blocked"
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert (
+            call_kwargs["properties"]["$ai_error_type"] == "output_guardrail_triggered"
+        )
+
+    def test_error_type_categorization_max_turns(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that MaxTurnsExceeded is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "MaxTurnsExceeded: Agent exceeded maximum turns"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "max_turns_exceeded"
+
+    def test_error_type_categorization_unknown(self, processor, mock_client, mock_span):
+        """Test that unknown errors are categorized as unknown."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "Some random error occurred"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "unknown"
+
+    def test_response_span_with_output_and_total_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test ResponseSpanData includes output choices and total tokens."""
+        # Create a mock response object
+        mock_response = MagicMock()
+        mock_response.id = "resp_123"
+        mock_response.model = "gpt-4o"
+        mock_response.output = [{"type": "message", "content": "Hello!"}]
+        mock_response.usage = MagicMock()
+        mock_response.usage.input_tokens = 25
+        mock_response.usage.output_tokens = 10
+
+        span_data = ResponseSpanData(
+            response=mock_response,
+            input="Hello, world!",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_generation"
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 35
+        assert call_kwargs["properties"]["$ai_output_choices"] == [
+            {"type": "message", "content": "Hello!"}
+        ]
+        assert call_kwargs["properties"]["$ai_response_id"] == "resp_123"
+
+    def test_speech_span_with_pass_through_properties(
+        self, processor, mock_client, mock_span
+    ):
+        """Test SpeechSpanData includes pass-through properties."""
+        span_data = SpeechSpanData(
+            input="Hello, how can I help you?",
+            output="base64_audio_data",
+            output_format="pcm",
+            model="tts-1",
+            model_config={"voice": "alloy", "speed": 1.0},
+            first_content_at="2024-01-01T00:00:00.500Z",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "speech"
+        assert call_kwargs["properties"]["$ai_model"] == "tts-1"
+        # Pass-through properties (no $ai_ prefix)
+        assert (
+            call_kwargs["properties"]["first_content_at"] == "2024-01-01T00:00:00.500Z"
+        )
+        assert call_kwargs["properties"]["audio_output_format"] == "pcm"
+        assert call_kwargs["properties"]["model_config"] == {
+            "voice": "alloy",
+            "speed": 1.0,
+        }
+        # Text input should be captured
+        assert call_kwargs["properties"]["$ai_input"] == "Hello, how can I help you?"
+
+    def test_transcription_span_with_pass_through_properties(
+        self, processor, mock_client, mock_span
+    ):
+        """Test TranscriptionSpanData includes pass-through properties."""
+        span_data = TranscriptionSpanData(
+            input="base64_audio_data",
+            input_format="pcm",
+            output="This is the transcribed text.",
+            model="whisper-1",
+            model_config={"language": "en"},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "transcription"
+        assert call_kwargs["properties"]["$ai_model"] == "whisper-1"
+        # Pass-through properties (no $ai_ prefix)
+        assert call_kwargs["properties"]["audio_input_format"] == "pcm"
+        assert call_kwargs["properties"]["model_config"] == {"language": "en"}
+        # Transcription output should be captured
+        assert (
+            call_kwargs["properties"]["$ai_output_state"]
+            == "This is the transcribed text."
+        )
+
+    def test_latency_calculation(self, processor, mock_client, mock_span):
+        """Test that latency is calculated correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        with patch("time.time") as mock_time:
+            mock_time.return_value = 1000.0
+            processor.on_span_start(mock_span)
+
+            mock_time.return_value = 1001.5  # 1.5 seconds later
+            processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_latency"] == pytest.approx(1.5, rel=0.01)
+
+    def test_groups_included_in_events(self, mock_client, mock_trace, mock_span):
+        """Test that groups are included in captured events."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            groups={"company": "acme", "team": "engineering"},
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["groups"] == {"company": "acme", "team": "engineering"}
+
+    def test_additional_properties_included(self, mock_client, mock_trace):
+        """Test that additional properties are included in events."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            properties={"environment": "production", "version": "1.0"},
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["environment"] == "production"
+        assert call_kwargs["properties"]["version"] == "1.0"
+
+    def test_shutdown_clears_state(self, processor):
+        """Test that shutdown clears internal state."""
+        processor._span_start_times["span_1"] = 1000.0
+        processor._trace_metadata["trace_1"] = {"name": "test"}
+
+        processor.shutdown()
+
+        assert len(processor._span_start_times) == 0
+        assert len(processor._trace_metadata) == 0
+
+    def test_force_flush_calls_client_flush(self, processor, mock_client):
+        """Test that force_flush calls client.flush()."""
+        processor.force_flush()
+        mock_client.flush.assert_called_once()
+
+    def test_generation_span_with_no_usage(self, processor, mock_client, mock_span):
+        """Test GenerationSpanData with no usage data defaults to zero tokens."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 0
+
+    def test_generation_span_with_partial_usage(
+        self, processor, mock_client, mock_span
+    ):
+        """Test GenerationSpanData with only input_tokens present."""
+        span_data = GenerationSpanData(
+            model="gpt-4o",
+            usage={"input_tokens": 42},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 42
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 42
+
+    def test_error_type_categorization_by_type_field_only(
+        self, processor, mock_client, mock_span
+    ):
+        """Test error categorization works when only the type field matches."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "Something went wrong",
+            "type": "ModelBehaviorError",
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
+
+    def test_distinct_id_resolved_from_trace_for_spans(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test that spans use the distinct_id resolved at trace start."""
+
+        def resolver(trace):
+            return f"user-{trace.name}"
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        # Start trace - this resolves and stores distinct_id
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        # End a span - should use the stored distinct_id from trace
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "user-Test Workflow"
+
+    def test_eviction_of_stale_entries(self, mock_client):
+        """Test that stale entries are evicted when max is exceeded."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+        )
+        processor._max_tracked_entries = 10
+
+        # Fill beyond max
+        for i in range(15):
+            processor._span_start_times[f"span_{i}"] = float(i)
+            processor._trace_metadata[f"trace_{i}"] = {"name": f"trace_{i}"}
+
+        processor._evict_stale_entries()
+
+        # Should have evicted half
+        assert len(processor._span_start_times) <= 10
+        assert len(processor._trace_metadata) <= 10
+
+
+class TestInstrumentHelper:
+    """Tests for the instrument() convenience function."""
+
+    def test_instrument_registers_processor(self, mock_client):
+        """Test that instrument() registers a processor."""
+        with patch("agents.tracing.add_trace_processor") as mock_add:
+            processor = instrument(
+                client=mock_client,
+                distinct_id="test-user",
+            )
+
+            mock_add.assert_called_once_with(processor)
+            assert isinstance(processor, PostHogTracingProcessor)
+
+    def test_instrument_with_privacy_mode(self, mock_client):
+        """Test instrument() respects privacy_mode."""
+        with patch("agents.tracing.add_trace_processor"):
+            processor = instrument(
+                client=mock_client,
+                privacy_mode=True,
+            )
+
+            assert processor._privacy_mode is True
+
+    def test_instrument_with_groups_and_properties(self, mock_client):
+        """Test instrument() accepts groups and properties."""
+        with patch("agents.tracing.add_trace_processor"):
+            processor = instrument(
+                client=mock_client,
+                groups={"company": "acme"},
+                properties={"env": "test"},
+            )
+
+            assert processor._groups == {"company": "acme"}
+            assert processor._properties == {"env": "test"}