AI wrappers: cancelled async calls (asyncio.CancelledError) are captured as successful $ai_generation (HTTP 200, no error) #1064

Summary

When a wrapped LLM call is cancelled before it returns (for example asyncio.wait_for times out, the request handler is cancelled, or a sibling task in a TaskGroup fails), the $ai_generation event is still captured, but it looks like a successful generation: $ai_http_status: 200, no $ai_is_error, and no tokens.

This skews LLM analytics. Cancelled calls are left out of the error rate, and their latency (often a whole timeout's worth) is counted as a successful call's latency. In our production data (Gemini, 30 days), these were about 0.4% of generations, they were among the slowest calls, and most of their latencies matched our timeout durations exactly.

Cause

call_llm_and_track_usage_async in posthog/ai/utils.py starts with http_status = 200, catches provider errors with except Exception, and captures the event in finally. asyncio.CancelledError subclasses BaseException, not Exception (since Python 3.8), so the except branch is skipped. The finally block then captures the event with the default 200 status and no error fields. Checked on 7.70.0 (latest at the time of writing).

Reproduction (posthog 7.70.0, Python 3.11)

import asyncio
from posthog import Posthog
from posthog.ai.utils import call_llm_and_track_usage_async

captured = []

def before_send(event):
    if event["event"] == "$ai_generation":
        captured.append(event["properties"])
    return None

ph = Posthog("phc_test", host="http://localhost:1", before_send=before_send, sync_mode=True)

async def slow_llm_call(**kwargs):
    await asyncio.sleep(10)

async def failing_llm_call(**kwargs):
    raise RuntimeError("provider error")

async def track(call):
    return await call_llm_and_track_usage_async(
        None, ph, "gemini", None, None, False, None, "", call, model="gemini-x"
    )

def show(label):
    p = captured[-1]
    print(label, {k: p.get(k) for k in ("$ai_is_error", "$ai_error", "$ai_http_status")})

async def main():
    try:
        await track(failing_llm_call)
    except RuntimeError:
        pass
    show("provider error      ->")
    try:
        await asyncio.wait_for(track(slow_llm_call), timeout=0.05)
    except asyncio.TimeoutError:
        pass
    show("cancelled (timeout) ->")

asyncio.run(main())

Output:

provider error      -> {'$ai_is_error': True, '$ai_error': 'provider error', '$ai_http_status': 0}
cancelled (timeout) -> {'$ai_is_error': None, '$ai_error': None, '$ai_http_status': 200}

Suggested fix

Handle asyncio.CancelledError explicitly: mark the event as an error, then re-raise so cancellation semantics are unchanged.

        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:
            ...

With this patch applied locally to 7.70.0, the reproduction prints {'$ai_is_error': True, '$ai_error': 'cancelled', '$ai_http_status': 0} for the cancelled call. The caller still gets TimeoutError from wait_for, and the provider-error event is unchanged.

The sync call_llm_and_track_usage has the same except Exception shape. There the analogous case is KeyboardInterrupt / SystemExit, which is probably lower priority. We have not checked the streaming code paths.

Workaround

Until this is fixed, we register a before_send hook on the client used by the AI wrappers. It marks $ai_generation events with $ai_http_status == 200, $ai_input_tokens == 0 and no output tokens as errors. That works for us, but it relies on inferring the shape of the event rather than knowing that the call was cancelled.