diff --git a/src/agents/models/_retry_runtime.py b/src/agents/models/_retry_runtime.py index 718931f403..5d6f427ad2 100644 --- a/src/agents/models/_retry_runtime.py +++ b/src/agents/models/_retry_runtime.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math import time from collections.abc import Iterator, Mapping from contextlib import contextmanager @@ -64,7 +65,7 @@ def parse_retry_after_ms(value: str | None) -> float | None: parsed = float(value) / 1000.0 except ValueError: return None - return parsed if parsed >= 0 else None + return parsed if math.isfinite(parsed) and parsed >= 0 else None def parse_retry_after_value(value: str | None) -> float | None: @@ -76,7 +77,7 @@ def parse_retry_after_value(value: str | None) -> float | None: except ValueError: parsed = None if parsed is not None: - return parsed if parsed >= 0 else None + return parsed if math.isfinite(parsed) and parsed >= 0 else None try: retry_datetime = parsedate_to_datetime(value) diff --git a/tests/models/test_retry_after_nonfinite.py b/tests/models/test_retry_after_nonfinite.py new file mode 100644 index 0000000000..c82fb6e6cf --- /dev/null +++ b/tests/models/test_retry_after_nonfinite.py @@ -0,0 +1,18 @@ +import pytest + +from agents.models._retry_runtime import parse_retry_after_ms, parse_retry_after_value + + +@pytest.mark.parametrize("value", ["Infinity", "+Infinity", "-Infinity", "NaN"]) +def test_retry_after_ms_rejects_non_finite_values(value: str) -> None: + assert parse_retry_after_ms(value) is None + + +@pytest.mark.parametrize("value", ["Infinity", "+Infinity", "-Infinity", "NaN"]) +def test_retry_after_rejects_non_finite_numeric_values(value: str) -> None: + assert parse_retry_after_value(value) is None + + +def test_retry_after_parsers_keep_finite_values() -> None: + assert parse_retry_after_ms("1500") == 1.5 + assert parse_retry_after_value("2.5") == 2.5