chore: import upstream snapshot with attribution

2026-07-13 12:39:17 +08:00
commit 4ed4e9ff99
1368 changed files with 334957 additions and 0 deletions
@@ -0,0 +1,570 @@
+from __future__ import annotations
+
+import pytest
+from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails
+from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
+
+from agents import Agent, Runner
+from agents.run_internal.agent_runner_helpers import snapshot_usage, usage_delta
+from agents.usage import (
+    RequestUsage,
+    Usage,
+    deserialize_usage,
+    model_usage_to_span_usage,
+    serialize_usage,
+)
+from tests.fake_model import FakeModel
+from tests.test_responses import get_text_message
+
+
+def test_usage_defaults_cache_write_tokens_to_zero() -> None:
+    usage = Usage()
+
+    assert usage.input_tokens_details.cached_tokens == 0
+    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 0
+
+
+@pytest.mark.asyncio
+async def test_runner_run_carries_request_usage_entries() -> None:
+    """Ensure usage produced by the model propagates to RunResult context."""
+    usage = Usage(
+        requests=1,
+        input_tokens=10,
+        output_tokens=5,
+        total_tokens=15,
+        request_usage_entries=[
+            RequestUsage(
+                input_tokens=10,
+                output_tokens=5,
+                total_tokens=15,
+                input_tokens_details=InputTokensDetails.model_validate(
+                    {"cache_write_tokens": 0, "cached_tokens": 0}
+                ),
+                output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+            )
+        ],
+    )
+    model = FakeModel(initial_output=[get_text_message("done")])
+    model.set_hardcoded_usage(usage)
+    agent = Agent(name="usage-agent", model=model)
+
+    result = await Runner.run(agent, input="hi")
+
+    propagated = result.context_wrapper.usage
+    assert propagated.requests == 1
+    assert propagated.total_tokens == 15
+    assert len(propagated.request_usage_entries) == 1
+    entry = propagated.request_usage_entries[0]
+    assert entry.input_tokens == 10
+    assert entry.output_tokens == 5
+    assert entry.total_tokens == 15
+
+
+def test_usage_add_aggregates_all_fields():
+    u1 = Usage(
+        requests=1,
+        input_tokens=10,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 5, "cached_tokens": 3}
+        ),
+        output_tokens=20,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
+        total_tokens=30,
+    )
+    u2 = Usage(
+        requests=2,
+        input_tokens=7,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 6, "cached_tokens": 4}
+        ),
+        output_tokens=8,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=6),
+        total_tokens=15,
+    )
+
+    u1.add(u2)
+
+    assert u1.requests == 3
+    assert u1.input_tokens == 17
+    assert u1.output_tokens == 28
+    assert u1.total_tokens == 45
+    assert u1.input_tokens_details.cached_tokens == 7
+    assert getattr(u1.input_tokens_details, "cache_write_tokens", None) == 11
+    assert u1.output_tokens_details.reasoning_tokens == 11
+
+
+def test_usage_add_aggregates_with_none_values():
+    u1 = Usage()
+    u2 = Usage(
+        requests=2,
+        input_tokens=7,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 4}
+        ),
+        output_tokens=8,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=6),
+        total_tokens=15,
+    )
+
+    u1.add(u2)
+
+    assert u1.requests == 2
+    assert u1.input_tokens == 7
+    assert u1.output_tokens == 8
+    assert u1.total_tokens == 15
+    assert u1.input_tokens_details.cached_tokens == 4
+    assert u1.output_tokens_details.reasoning_tokens == 6
+
+
+def test_request_usage_creation():
+    """Test that RequestUsage is created correctly."""
+    request_usage = RequestUsage(
+        input_tokens=100,
+        output_tokens=200,
+        total_tokens=300,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 10}
+        ),
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
+    )
+
+    assert request_usage.input_tokens == 100
+    assert request_usage.output_tokens == 200
+    assert request_usage.total_tokens == 300
+    assert request_usage.input_tokens_details.cached_tokens == 10
+    assert request_usage.output_tokens_details.reasoning_tokens == 20
+
+
+def test_usage_add_preserves_single_request():
+    """Test that adding a single request Usage creates an RequestUsage entry."""
+    u1 = Usage()
+    u2 = Usage(
+        requests=1,
+        input_tokens=100,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 10}
+        ),
+        output_tokens=200,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
+        total_tokens=300,
+    )
+
+    u1.add(u2)
+
+    # Should preserve the request usage details
+    assert len(u1.request_usage_entries) == 1
+    request_usage = u1.request_usage_entries[0]
+    assert request_usage.input_tokens == 100
+    assert request_usage.output_tokens == 200
+    assert request_usage.total_tokens == 300
+    assert request_usage.input_tokens_details.cached_tokens == 10
+    assert request_usage.output_tokens_details.reasoning_tokens == 20
+
+
+def test_usage_add_ignores_zero_token_requests():
+    """Test that zero-token requests don't create request_usage_entries."""
+    u1 = Usage()
+    u2 = Usage(
+        requests=1,
+        input_tokens=0,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 0}
+        ),
+        output_tokens=0,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+        total_tokens=0,
+    )
+
+    u1.add(u2)
+
+    # Should not create a request_usage_entry for zero tokens
+    assert len(u1.request_usage_entries) == 0
+
+
+def test_usage_add_ignores_multi_request_usage():
+    """Test that multi-request Usage objects don't create request_usage_entries."""
+    u1 = Usage()
+    u2 = Usage(
+        requests=3,  # Multiple requests
+        input_tokens=100,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 10}
+        ),
+        output_tokens=200,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
+        total_tokens=300,
+    )
+
+    u1.add(u2)
+
+    # Should not create a request usage entry for multi-request usage
+    assert len(u1.request_usage_entries) == 0
+
+
+def test_usage_add_merges_existing_request_usage_entries():
+    """Test that existing request_usage_entries are merged when adding Usage objects."""
+    # Create first usage with request_usage_entries
+    u1 = Usage()
+    u2 = Usage(
+        requests=1,
+        input_tokens=100,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 10}
+        ),
+        output_tokens=200,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
+        total_tokens=300,
+    )
+    u1.add(u2)
+
+    # Create second usage with request_usage_entries
+    u3 = Usage(
+        requests=1,
+        input_tokens=50,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 5}
+        ),
+        output_tokens=75,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
+        total_tokens=125,
+    )
+
+    u1.add(u3)
+
+    # Should have both request_usage_entries
+    assert len(u1.request_usage_entries) == 2
+
+    # First request
+    first = u1.request_usage_entries[0]
+    assert first.input_tokens == 100
+    assert first.output_tokens == 200
+    assert first.total_tokens == 300
+
+    # Second request
+    second = u1.request_usage_entries[1]
+    assert second.input_tokens == 50
+    assert second.output_tokens == 75
+    assert second.total_tokens == 125
+
+
+def test_usage_add_with_pre_existing_request_usage_entries():
+    """Test adding Usage objects that already have request_usage_entries."""
+    u1 = Usage()
+
+    # Create a usage with request_usage_entries
+    u2 = Usage(
+        requests=1,
+        input_tokens=100,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 10}
+        ),
+        output_tokens=200,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
+        total_tokens=300,
+    )
+    u1.add(u2)
+
+    # Create another usage with request_usage_entries
+    u3 = Usage(
+        requests=1,
+        input_tokens=50,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 5}
+        ),
+        output_tokens=75,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
+        total_tokens=125,
+    )
+
+    # Add u3 to u1
+    u1.add(u3)
+
+    # Should have both request_usage_entries
+    assert len(u1.request_usage_entries) == 2
+    assert u1.request_usage_entries[0].input_tokens == 100
+    assert u1.request_usage_entries[1].input_tokens == 50
+
+
+def test_usage_add_preserves_existing_entries_when_top_level_also_set():
+    """When `other` has both top-level single-request fields AND pre-populated
+    `request_usage_entries`, the existing entries (which carry the authoritative
+    nested token details) must not be discarded in favor of a synthesized entry
+    built from only the top-level fields.
+    """
+    u1 = Usage()
+    u2 = Usage(
+        requests=1,
+        input_tokens=100,
+        output_tokens=50,
+        total_tokens=150,
+        request_usage_entries=[
+            RequestUsage(
+                input_tokens=100,
+                output_tokens=50,
+                total_tokens=150,
+                input_tokens_details=InputTokensDetails.model_validate(
+                    {"cache_write_tokens": 0, "cached_tokens": 10}
+                ),
+                output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
+            )
+        ],
+    )
+
+    u1.add(u2)
+
+    # The pre-populated entry must be preserved — including its nested details —
+    # rather than being replaced by a synthesized entry with zeroed-out details.
+    assert len(u1.request_usage_entries) == 1
+    entry = u1.request_usage_entries[0]
+    assert entry.input_tokens_details.cached_tokens == 10
+    assert entry.output_tokens_details.reasoning_tokens == 5
+
+
+def test_usage_request_usage_entries_default_empty():
+    """Test that request_usage_entries defaults to an empty list."""
+    u = Usage()
+    assert u.request_usage_entries == []
+
+
+def test_anthropic_cost_calculation_scenario():
+    """Test a realistic scenario for Sonnet 4.5 cost calculation with 200K token thresholds."""
+    # Simulate 3 API calls: 100K, 150K, and 80K input tokens each
+    # None exceed 200K, so they should all use the lower pricing tier
+
+    usage = Usage()
+
+    # First request: 100K input tokens
+    req1 = Usage(
+        requests=1,
+        input_tokens=100_000,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 0}
+        ),
+        output_tokens=50_000,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+        total_tokens=150_000,
+    )
+    usage.add(req1)
+
+    # Second request: 150K input tokens
+    req2 = Usage(
+        requests=1,
+        input_tokens=150_000,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 0}
+        ),
+        output_tokens=75_000,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+        total_tokens=225_000,
+    )
+    usage.add(req2)
+
+    # Third request: 80K input tokens
+    req3 = Usage(
+        requests=1,
+        input_tokens=80_000,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 0, "cached_tokens": 0}
+        ),
+        output_tokens=40_000,
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+        total_tokens=120_000,
+    )
+    usage.add(req3)
+
+    # Verify aggregated totals
+    assert usage.requests == 3
+    assert usage.input_tokens == 330_000  # 100K + 150K + 80K
+    assert usage.output_tokens == 165_000  # 50K + 75K + 40K
+    assert usage.total_tokens == 495_000  # 150K + 225K + 120K
+
+    # Verify request_usage_entries preservation
+    assert len(usage.request_usage_entries) == 3
+    assert usage.request_usage_entries[0].input_tokens == 100_000
+    assert usage.request_usage_entries[1].input_tokens == 150_000
+    assert usage.request_usage_entries[2].input_tokens == 80_000
+
+    # All request_usage_entries are under 200K threshold
+    for req in usage.request_usage_entries:
+        assert req.input_tokens < 200_000
+        assert req.output_tokens < 200_000
+
+
+def test_usage_normalizes_none_token_details():
+    # Some providers don't populate optional token detail fields
+    # (cached_tokens, reasoning_tokens), and the OpenAI SDK's generated
+    # code can bypass Pydantic validation (e.g., via model_construct),
+    # allowing None values. We normalize these to 0 to prevent TypeErrors.
+
+    # Test entire objects being None (BeforeValidator)
+    usage = Usage(
+        requests=1,
+        input_tokens=100,
+        input_tokens_details=None,  # type: ignore[arg-type]
+        output_tokens=50,
+        output_tokens_details=None,  # type: ignore[arg-type]
+        total_tokens=150,
+    )
+    assert usage.input_tokens_details.cached_tokens == 0
+    assert usage.output_tokens_details.reasoning_tokens == 0
+
+    # Test fields within objects being None (__post_init__)
+    input_details = InputTokensDetails.model_validate({"cache_write_tokens": 0, "cached_tokens": 0})
+    input_details.__dict__["cached_tokens"] = None
+    input_details.__dict__["cache_write_tokens"] = None
+
+    output_details = OutputTokensDetails(reasoning_tokens=0)
+    output_details.__dict__["reasoning_tokens"] = None
+
+    usage = Usage(
+        requests=1,
+        input_tokens=100,
+        input_tokens_details=input_details,
+        output_tokens=50,
+        output_tokens_details=output_details,
+        total_tokens=150,
+    )
+
+    # __post_init__ should normalize None to 0
+    assert usage.input_tokens_details.cached_tokens == 0
+    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 0
+    assert usage.output_tokens_details.reasoning_tokens == 0
+
+
+def test_usage_normalizes_chat_completions_types():
+    # Chat Completions API uses PromptTokensDetails and CompletionTokensDetails,
+    # while Usage expects InputTokensDetails and OutputTokensDetails (Responses API).
+    # The BeforeValidator should convert between these types.
+
+    prompt_details = PromptTokensDetails.model_validate(
+        {
+            "audio_tokens": 10,
+            "cached_tokens": 50,
+            "cache_write_tokens": 7,
+        }
+    )
+    completion_details = CompletionTokensDetails(
+        accepted_prediction_tokens=5,
+        audio_tokens=10,
+        reasoning_tokens=100,
+        rejected_prediction_tokens=2,
+    )
+
+    usage = Usage(
+        requests=1,
+        input_tokens=200,
+        input_tokens_details=prompt_details,  # type: ignore[arg-type]
+        output_tokens=150,
+        output_tokens_details=completion_details,  # type: ignore[arg-type]
+        total_tokens=350,
+    )
+
+    # Should convert to Responses API types, extracting the relevant fields
+    assert isinstance(usage.input_tokens_details, InputTokensDetails)
+    assert usage.input_tokens_details.cached_tokens == 50
+    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 7
+
+    assert isinstance(usage.output_tokens_details, OutputTokensDetails)
+    assert usage.output_tokens_details.reasoning_tokens == 100
+
+
+def test_usage_serialization_preserves_cache_write_tokens() -> None:
+    usage = Usage(
+        requests=1,
+        input_tokens=20,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 7, "cached_tokens": 3}
+        ),
+        output_tokens=5,
+        total_tokens=25,
+        request_usage_entries=[
+            RequestUsage(
+                input_tokens=20,
+                output_tokens=5,
+                total_tokens=25,
+                input_tokens_details=InputTokensDetails.model_validate(
+                    {"cache_write_tokens": 7, "cached_tokens": 3}
+                ),
+                output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+            )
+        ],
+    )
+
+    serialized = serialize_usage(usage)
+    restored = deserialize_usage(serialized)
+
+    assert serialized["input_tokens_details"] == [{"cached_tokens": 3, "cache_write_tokens": 7}]
+    assert getattr(restored.input_tokens_details, "cache_write_tokens", None) == 7
+    assert (
+        getattr(
+            restored.request_usage_entries[0].input_tokens_details,
+            "cache_write_tokens",
+            None,
+        )
+        == 7
+    )
+
+
+def test_usage_deserialization_defaults_legacy_cache_write_tokens() -> None:
+    restored = deserialize_usage(
+        {
+            "requests": 1,
+            "input_tokens": 20,
+            "output_tokens": 5,
+            "total_tokens": 25,
+            "input_tokens_details": [{"cached_tokens": 3}],
+            "request_usage_entries": [
+                {
+                    "input_tokens": 20,
+                    "output_tokens": 5,
+                    "total_tokens": 25,
+                    "input_tokens_details": {"cached_tokens": 3},
+                }
+            ],
+        }
+    )
+
+    assert restored.input_tokens_details.cached_tokens == 3
+    assert getattr(restored.input_tokens_details, "cache_write_tokens", None) == 0
+    assert restored.request_usage_entries[0].input_tokens_details.cached_tokens == 3
+    assert (
+        getattr(
+            restored.request_usage_entries[0].input_tokens_details,
+            "cache_write_tokens",
+            None,
+        )
+        == 0
+    )
+
+
+def test_usage_snapshot_delta_and_span_preserve_cache_write_tokens() -> None:
+    start = Usage(
+        requests=1,
+        input_tokens=10,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 2, "cached_tokens": 3}
+        ),
+        output_tokens=4,
+        total_tokens=14,
+    )
+    end = Usage(
+        requests=2,
+        input_tokens=30,
+        input_tokens_details=InputTokensDetails.model_validate(
+            {"cache_write_tokens": 9, "cached_tokens": 8}
+        ),
+        output_tokens=10,
+        total_tokens=40,
+    )
+
+    snapshot = snapshot_usage(start)
+    delta = usage_delta(snapshot, end)
+    span_usage = model_usage_to_span_usage(delta)
+
+    assert getattr(snapshot.input_tokens_details, "cache_write_tokens", None) == 2
+    assert delta.input_tokens_details.cached_tokens == 5
+    assert getattr(delta.input_tokens_details, "cache_write_tokens", None) == 7
+    assert span_usage["input_tokens_details"] == {
+        "cached_tokens": 5,
+        "cache_write_tokens": 7,
+    }