Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions backend/app/core/retry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import httpx
import requests


def is_transient(error: BaseException) -> bool:
"""5xx from upstream, or a connection that never reached it."""
if isinstance(error, (httpx.HTTPStatusError, requests.HTTPError)):
return error.response is not None and error.response.status_code >= 500
return isinstance(
error, (httpx.ConnectError, httpx.ConnectTimeout, requests.ConnectionError)
)
50 changes: 37 additions & 13 deletions backend/app/services/llm/guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@
import httpx
from fastapi import HTTPException
from opentelemetry import trace
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential

from app.core.config import settings
from app.core.retry import is_transient
from app.models.llm.request import InlineValidator, InlineValidatorTypeEnum, Validator

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -91,7 +93,13 @@ def proxy_guardrails_request(
),
httpx.Client(timeout=GUARDRAILS_PROXY_TIMEOUT_SECONDS) as client,
):
response = client.request(
# Status isn't raised, so only connect failures retry: safe for any method.
response = retry(
stop=stop_after_attempt(4),
wait=wait_exponential(max=10),
retry=retry_if_exception(is_transient),
reraise=True,
)(client.request)(
method, url, params=query, json=json_body, headers=headers
)
except httpx.RequestError as e:
Expand Down Expand Up @@ -371,19 +379,29 @@ def run_guardrails_validation(
),
httpx.Client(timeout=45.0) as client,
):
response = client.post(
url,
json=payload,
params={"suppress_pass_logs": str(suppress_pass_logs).lower()},
headers=headers,
)
elapsed_ms = int((time.monotonic() - started) * 1000)
logger.info(
f"[run_guardrails_validation] Response received | job_id={job_id}, "
f"status={response.status_code}, elapsed_ms={elapsed_ms}, "

@retry(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can combine this into one single decorator I guess, instead of repeating everywhere

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a decorator. So thought of keeping it staright forward

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That was Claude's suggestion. But overrode it to make the retry straight forward

stop=stop_after_attempt(4),
wait=wait_exponential(max=10),
retry=retry_if_exception(is_transient),
reraise=True,
)
response.raise_for_status()
return response.json()
def _post() -> dict[str, Any]:
response = client.post(
url,
json=payload,
params={"suppress_pass_logs": str(suppress_pass_logs).lower()},
headers=headers,
)
elapsed_ms = int((time.monotonic() - started) * 1000)
logger.info(
f"[run_guardrails_validation] Response received | job_id={job_id}, "
f"status={response.status_code}, elapsed_ms={elapsed_ms}, "
)
response.raise_for_status()
return response.json()

return _post()
except Exception as e:
elapsed_ms = int((time.monotonic() - started) * 1000)
if _is_auth_error(e):
Expand Down Expand Up @@ -458,6 +476,12 @@ def _build_params(validator_ids: list[UUID]) -> dict[str, Any]:
httpx.Client(timeout=10.0) as client,
):

@retry(
stop=stop_after_attempt(4),
wait=wait_exponential(max=10),
retry=retry_if_exception(is_transient),
reraise=True,
)
def _fetch_by_ids(validator_ids: list[UUID]) -> list[dict[str, Any]]:
if not validator_ids:
return []
Expand Down
48 changes: 34 additions & 14 deletions backend/app/services/llm/jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,12 @@
from gevent import Timeout
from opentelemetry import trace
from sqlmodel import Session
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential

from app.celery.utils import start_llm_chain_job, start_llm_job
from app.core.db import engine
from app.models.provider_credentials import Provider
from app.core.providers import Provider
from app.core.retry import is_transient
from app.core.langfuse.langfuse import observe_llm_execution
from app.core.telemetry import (
LogModuleEnum,
Expand Down Expand Up @@ -860,6 +862,28 @@ def _run_output_guardrails(
return result


# ReadTimeout isn't retried: the proxy may already have run (and billed) the call.
@retry(
stop=stop_after_attempt(4),
wait=wait_exponential(max=10),
retry=retry_if_exception(is_transient),
reraise=True,
)
def _post_to_proxy(url: str, api_key: str, input_text: str) -> dict:
with httpx.Client(timeout=60.0) as client:
resp = client.post(
url,
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"Accept": "application/json",
},
json={"input": input_text},
)
resp.raise_for_status()
return resp.json()


def _execute_proxy_call(
session: Session,
*,
Expand Down Expand Up @@ -919,6 +943,10 @@ def _execute_proxy_call(
)
return BlockResult(error=f"Failed to create LLM call record: {str(e)}")

# End the read transaction so no pooled connection sits idle during the proxy
# request and its retry backoff; nothing is pending, so this only releases it.
session.commit()

ctx.metrics_started(provider=Provider.PROXY.value, model="")
proxy_data: dict | None = None

Expand All @@ -928,18 +956,9 @@ def _execute_proxy_call(

try:
with suppress_http_instrumentation():
with httpx.Client(timeout=60.0) as client:
resp = client.post(
client_llm_url,
headers={
"Authorization": f"Bearer {creds['api_key']}",
"Content-Type": "application/json",
"Accept": "application/json",
},
json={"input": query.input.content.value},
)
resp.raise_for_status()
proxy_data = resp.json()
proxy_data = _post_to_proxy(
client_llm_url, creds["api_key"], query.input.content.value
)
except httpx.HTTPError as e:
proxy_span.record_exception(e)
proxy_span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
Expand All @@ -951,7 +970,8 @@ def _execute_proxy_call(
return BlockResult(
error=f"Proxy call failed: {str(e)}",
llm_call_id=ctx.llm_call_id,
retryable=True,
# Transient failures were already retried; don't stack fast-eval retries.
retryable=not is_transient(e),
)

try:
Expand Down
6 changes: 6 additions & 0 deletions backend/app/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,12 @@ def aws_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AWS_DEFAULT_REGION", settings.AWS_DEFAULT_REGION)


@pytest.fixture
def no_retry_backoff(monkeypatch: pytest.MonkeyPatch) -> None:
"""Tenacity binds nap.sleep at import (incl. nested decorators), so patch time.sleep."""
monkeypatch.setattr("tenacity.nap.time.sleep", lambda _seconds: None)


@pytest.fixture(scope="function")
def client(db: Session) -> Generator[TestClient, None, None]:
app.dependency_overrides[get_db] = lambda: db
Expand Down
67 changes: 62 additions & 5 deletions backend/app/tests/core/test_callback_ssrf.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import itertools
import json
import socket
from typing import Any
Expand All @@ -16,6 +17,8 @@
sign_webhook_payload,
)

pytestmark = pytest.mark.usefixtures("no_retry_backoff")


class TestIsPrivateIP:
"""Test suite for _is_private_ip function."""
Expand Down Expand Up @@ -237,7 +240,7 @@ def test_successful_callback(
"""Test successful callback execution."""
mock_session = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.status_code = 200
mock_response.iter_content.return_value = [b"test"]
mock_session.post.return_value = mock_response
mock_session_class.return_value.__enter__.return_value = mock_session
Expand Down Expand Up @@ -272,7 +275,7 @@ def test_callback_http_error(
"""Test that callback returns False on HTTP errors."""
mock_session = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = requests.HTTPError("404 Not Found")
mock_response.status_code = 404
mock_session.post.return_value = mock_response
mock_session_class.return_value.__enter__.return_value = mock_session

Expand All @@ -288,7 +291,7 @@ def test_callback_disables_redirects(
"""Test that redirects are disabled to prevent redirect-based SSRF."""
mock_session = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.status_code = 200
mock_response.iter_content.return_value = [b"test"]
mock_session.post.return_value = mock_response
mock_session_class.return_value.__enter__.return_value = mock_session
Expand All @@ -306,7 +309,7 @@ def test_callback_uses_timeout(
"""Test that callback uses configured timeouts."""
mock_session = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.status_code = 200
mock_response.iter_content.return_value = [b"test"]
mock_session.post.return_value = mock_response
mock_session_class.return_value.__enter__.return_value = mock_session
Expand All @@ -326,7 +329,7 @@ def test_callback_sends_json_data(
"""Test that callback sends data as JSON."""
mock_session = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.status_code = 200
mock_response.iter_content.return_value = [b"test"]
mock_session.post.return_value = mock_response
mock_session_class.return_value.__enter__.return_value = mock_session
Expand Down Expand Up @@ -378,3 +381,57 @@ def test_signature_matches_manual_hmac(self):
sig, returned_ts = sign_webhook_payload(secret, body, timestamp_ms=ts)
assert sig == expected
assert returned_ts == ts


def _callback_response(status_code: int) -> requests.Response:
response = requests.Response()
response.status_code = status_code
return response


@patch("app.utils.validate_callback_url")
class TestSendCallbackRetry:
URL = "https://api.example.com/callback"

def test_recovers_after_transient_failures(self, _validate: Any) -> None:
# 302 counts as a failure: redirects are off, so nothing was delivered.
outcomes = [
_callback_response(302),
requests.ConnectionError("reset"),
_callback_response(200),
]
with patch.object(requests.Session, "post", side_effect=outcomes) as post:
assert send_callback(self.URL, {"status": "done"}) is True
assert post.call_count == 3

@pytest.mark.parametrize("status_code", [503, 404])
def test_gives_up_after_four_attempts(
self, _validate: Any, status_code: int
) -> None:
with patch.object(
requests.Session, "post", return_value=_callback_response(status_code)
) as post:
assert send_callback(self.URL, {"status": "done"}) is False
assert post.call_count == 4

def test_each_attempt_is_signed_with_a_fresh_timestamp(
self, _validate: Any
) -> None:
with (
patch("app.utils.time.time", side_effect=itertools.count(1_700_000_000)),
patch.object(
requests.Session, "post", return_value=_callback_response(503)
) as post,
):
send_callback(self.URL, {"status": "done"}, webhook_secret="s3cret")

timestamps = {
c.kwargs["headers"]["X-Webhook-Timestamp"] for c in post.mock_calls
}
assert len(timestamps) == 4

def test_invalid_url_is_never_posted(self, validate: Any) -> None:
validate.side_effect = ValueError("private IP")
with patch.object(requests.Session, "post") as post:
assert send_callback(self.URL, {"status": "done"}) is False
post.assert_not_called()
27 changes: 27 additions & 0 deletions backend/app/tests/core/test_retry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
import httpx
import pytest

from app.core.retry import is_transient

_REQUEST = httpx.Request("POST", "https://upstream.example/v1")


def _status_error(code: int) -> httpx.HTTPStatusError:
return httpx.HTTPStatusError(
"upstream", request=_REQUEST, response=httpx.Response(code, request=_REQUEST)
)


@pytest.mark.parametrize(
"error, expected",
[
(_status_error(500), True),
(_status_error(404), False),
(httpx.ConnectError("refused"), True),
(httpx.ConnectTimeout("slow handshake"), True),
# The request may have reached upstream, so a retry could double-execute.
(httpx.ReadTimeout("no response"), False),
],
)
def test_is_transient(error: BaseException, expected: bool) -> None:
assert is_transient(error) is expected
13 changes: 12 additions & 1 deletion backend/app/tests/services/llm/test_execute_llm_call.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
from app.services.llm.jobs import execute_llm_call
from app.tests.utils.utils import get_project

pytestmark = pytest.mark.usefixtures("no_retry_backoff")

VALIDATOR_CONFIG_ID = "00000000-0000-0000-0000-000000000001"
PROXY_URL = "https://api.tap.example/v1/predictions"

Expand Down Expand Up @@ -381,13 +383,22 @@ def test_provider_failure_is_retryable(self, db: Session, provider):

assert result.retryable is True

def test_proxy_transport_failure_is_retryable(
def test_proxy_connect_failure_is_not_retryable_after_inline_retries(
self, db: Session, provider, proxy_http
):
proxy_http.post.side_effect = httpx.ConnectError("connection refused")

result = call(db, build_blob(PROXY_COMPLETION))

assert result.error.startswith("Proxy call failed:")
assert result.retryable is False

def test_proxy_4xx_is_retryable(self, db: Session, provider, proxy_http):
request = httpx.Request("POST", PROXY_URL)
proxy_http.post.return_value = httpx.Response(400, request=request)

result = call(db, build_blob(PROXY_COMPLETION))

assert result.error.startswith("Proxy call failed:")
assert result.retryable is True

Expand Down
Loading
Loading