diff --git a/.sampo/changesets/venerable-baroness-lemminkainen.md b/.sampo/changesets/venerable-baroness-lemminkainen.md new file mode 100644 index 000000000..5dc5c3502 --- /dev/null +++ b/.sampo/changesets/venerable-baroness-lemminkainen.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Mark cancelled async LLM calls as errors while preserving cancellation. diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index ffc2c864c..c943bcc28 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -1,3 +1,4 @@ +import asyncio import logging import time import uuid @@ -648,6 +649,13 @@ async def call_llm_and_track_usage_async( try: response = await call_async_method(**kwargs) + except asyncio.CancelledError: + http_status = 0 + error_params = { + "$ai_is_error": True, + "$ai_error": "cancelled", + } + raise except Exception as exc: error = exc http_status = getattr( diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 8dc51ca34..b4e63c019 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -1,3 +1,4 @@ +import asyncio import os from unittest.mock import AsyncMock, MagicMock, patch @@ -150,6 +151,32 @@ async def test_async_client_basic_generation( assert props["$ai_latency"] > 0 +async def test_async_client_timeout_captures_error( + mock_client, mock_google_genai_client +): + async def slow_call(**kwargs): + await asyncio.Event().wait() + + mock_google_genai_client.aio.models.generate_content = AsyncMock( + side_effect=slow_call + ) + client = AsyncClient(api_key="test-key", posthog_client=mock_client) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + client.models.generate_content(model="gemini-x", contents="Hello"), + timeout=0.01, + ) + + assert mock_client.capture.call_count == 1 + event = mock_client.capture.call_args.kwargs + assert event["event"] == "$ai_generation" + props = event["properties"] + assert props["$ai_http_status"] == 0 + assert props["$ai_is_error"] is True + assert props["$ai_error"] == "cancelled" + + async def test_async_client_streaming_with_generate_content_stream( mock_client, mock_google_genai_client ): diff --git a/posthog/test/ai/test_async_cancellation.py b/posthog/test/ai/test_async_cancellation.py new file mode 100644 index 000000000..f247bbcdf --- /dev/null +++ b/posthog/test/ai/test_async_cancellation.py @@ -0,0 +1,87 @@ +import asyncio + +import pytest + +from posthog import Posthog +from posthog.ai.utils import call_llm_and_track_usage_async + + +@pytest.fixture +def captured_generations(): + events = [] + + def before_send(event): + if event["event"] == "$ai_generation": + events.append(event) + return None + + client = Posthog("phc_test", before_send=before_send, sync_mode=True) + try: + yield client, events + finally: + client.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["gemini", "openai", "anthropic"]) +async def test_cancelled_call_captures_error_and_preserves_cancellation( + captured_generations, provider +): + client, events = captured_generations + started = asyncio.Event() + + async def slow_call(**kwargs): + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task( + call_llm_and_track_usage_async( + "user-1", + client, + provider, + "trace-1", + None, + False, + None, + "", + slow_call, + model="test-model", + ) + ) + await started.wait() + task.cancel("request cancelled") + + with pytest.raises(asyncio.CancelledError): + await task + + assert task.cancelled() + assert len(events) == 1 + props = events[0]["properties"] + assert props["$ai_http_status"] == 0 + assert props["$ai_is_error"] is True + assert props["$ai_error"] == "cancelled" + assert props["$ai_provider"] == provider + assert props["$ai_trace_id"] == "trace-1" + assert "$ai_input_tokens" not in props + assert "$ai_output_tokens" not in props + + +@pytest.mark.asyncio +async def test_provider_error_is_captured_and_reraised(captured_generations): + client, events = captured_generations + error = RuntimeError("provider error") + + async def failing_call(**kwargs): + raise error + + with pytest.raises(RuntimeError) as exc_info: + await call_llm_and_track_usage_async( + None, client, "gemini", None, None, False, None, "", failing_call + ) + + assert exc_info.value is error + assert len(events) == 1 + props = events[0]["properties"] + assert props["$ai_http_status"] == 0 + assert props["$ai_is_error"] is True + assert props["$ai_error"] == "provider error" diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index ac29c3a9f..1d5afe45c 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -152,7 +152,6 @@ def echo(context: str, conversation_id: str, llm_model: str) -> str: tool = Tool.from_function(echo) if source == "middleware": child.add_middleware(ToolInjectionMiddleware(tools=[tool])) - assert await child.get_tool("echo") is None else: child.add_tool(tool) server = FastMCP("example-app") if source == "mounted" else child @@ -225,7 +224,6 @@ def echo(text: str) -> str: server = FastMCP("example-middleware-tools") server.add_middleware(ToolInjectionMiddleware(tools=[Tool.from_function(echo)])) - assert await server.get_tool("echo") is None sink = FakeClient() instrument( server,