video-image-production-build/tests/observability/test_context_tracking.py

from __future__ import annotations

import httpx

from production_build_agents.observability.context_tracking import (
    activate_context_tracking,
    begin_model_call,
    context_tracking_active,
    deactivate_context_tracking,
    finish_model_call,
    observe_transport_request,
    observe_transport_response,
    observed_http_clients,
)


def _snapshot() -> dict:
    return {
        "physical_call_id": "physical-1",
        "logical_call_index": 1,
        "logical_call_key": "logical-1",
        "identity": {"agent_run_id": "agent-1"},
        "previous": None,
    }


def test_transport_attempts_and_retries_are_recorded() -> None:
    capture_token = activate_context_tracking()
    call_token = begin_model_call(_snapshot())
    try:
        first = httpx.Request(
            "POST",
            "https://provider.example/v1/chat",
            headers={"x-stainless-retry-count": "0"},
        )
        observe_transport_request(first)
        observe_transport_response(httpx.Response(429, request=first))
        second = httpx.Request(
            "POST",
            "https://provider.example/v1/chat",
            headers={"x-stainless-retry-count": "1"},
        )
        observe_transport_request(second)
        observe_transport_response(httpx.Response(200, request=second))

        terminal = finish_model_call(
            call_token,
            message=None,
            response=None,
            status="success",
        )
    finally:
        deactivate_context_tracking(capture_token)

    transport = terminal["transport"]
    assert transport["observed"] is True
    assert transport["attempt_count"] == 2
    assert transport["retry_count"] == 1
    assert [item["status_code"] for item in transport["attempts"]] == [
        429,
        200,
    ]
    assert "provider.example" not in str(transport)


def test_unobserved_transport_is_unknown_not_zero() -> None:
    call_token = begin_model_call(_snapshot())
    terminal = finish_model_call(
        call_token,
        message=None,
        response=None,
        status="failure",
        error_type="TimeoutError",
    )

    assert terminal["transport"] == {
        "observed": False,
        "attempt_count": None,
        "retry_count": None,
        "attempts": [],
    }
    assert terminal["error_type"] == "TimeoutError"


def test_transport_exception_marks_attempt_without_response_as_error() -> None:
    capture_token = activate_context_tracking()
    call_token = begin_model_call(_snapshot())
    try:
        request = httpx.Request(
            "POST",
            "https://provider.example/v1/chat",
            headers={"x-stainless-retry-count": "0"},
        )
        observe_transport_request(request)
        terminal = finish_model_call(
            call_token,
            message=None,
            response=None,
            status="failure",
            error_type="ConnectError",
        )
    finally:
        deactivate_context_tracking(capture_token)

    assert terminal["transport"]["attempt_count"] == 1
    assert terminal["transport"]["retry_count"] == 0
    assert terminal["transport"]["attempts"][0]["outcome"] == "error"
    assert terminal["transport"]["attempts"][0]["status_code"] is None


def test_http_client_is_only_enabled_inside_full_capture() -> None:
    assert context_tracking_active() is False
    assert observed_http_clients("openai") == {}

    token = activate_context_tracking()
    try:
        assert context_tracking_active() is True
        assert "http_client" in observed_http_clients("openai")
        assert observed_http_clients("anthropic") == {}
    finally:
        deactivate_context_tracking(token)