fix(chat): honor provider Retry-After headers on rate limits

[Yigtwxx]

RetryAfterMiddleware was written to obey a provider's Retry-After header instead of guessing with exponential backoff. The branch that reads the header is gated on isinstance(headers, dict), but the header object litellm actually produces is httpx.Headers — a Mapping that is not a dict subclass — so that branch has never executed in production.

Description

Two changes in app/agents/chat/shared/middleware/retry_after.py:

1. _extract_retry_after_seconds now tests isinstance(headers, Mapping) instead of isinstance(headers, dict).
2. _delay_for_attempt caps the resulting delay at max_delay, which previously bounded only the exponential backoff.

Symptom

A provider returns 429 with Retry-After: 45. SurfSense ignores it, sleeps its own 1s / 2s / 4s backoff, burns all three retries inside roughly seven seconds, and fails the turn — the exact behaviour the module docstring says it was written to replace.

Root cause

_extract_retry_after_seconds reads exc.response.headers and then gates on isinstance(headers, dict). litellm.exceptions.RateLimitError.__init__ rebuilds the error's response unconditionally:

self.response = httpx.Response(
    status_code=429,
    headers=_response_headers,
    request=httpx.Request(method="POST", url=" https://cloud.google.com/vertex-ai/"),
)

so .headers is always httpx.Headers, which subclasses MutableMapping, not dict. Measured against the versions this repo pins:

litellm 1.88.1 | openai 2.24.0 | httpx 0.28.1
RateLimitError mro: ['RateLimitError', 'RateLimitError', 'APIStatusError', 'APIError', 'OpenAIError', 'Exception', 'BaseException']
exc.response: Response | headers: Headers
isinstance(headers, dict)    -> False
isinstance(headers, Mapping) -> True

The models are driven through ChatLiteLLM (app/agents/chat/runtime/llm_config.py, app/services/llm_service.py), so this is the only shape the middleware ever sees. With the branch skipped, the function falls through to its message regex, which does not match litellm's message ("litellm.RateLimitError: ..."), and returns None.

The existing tests did not catch this because their fixture builds the headers as a plain dict[str, str] (tests/unit/agents/new_chat/test_retry_after.py), so they exercise a branch production never reaches.

The second defect, and why it belongs in the same PR

_delay_for_attempt did return max(backoff, header). max_delay is documented as "Cap on per-attempt delay in seconds" but only ever constrained backoff. That was harmless while header was permanently 0.0; fixing the first defect makes it reachable. A retry-after-ms: 3600000 from a misconfigured gateway would become await asyncio.sleep(3600), and this loop runs inside the live chat turn — it holds the SSE stream, the thread's busy-mutex lock and the DB session for the whole sleep. The delay is now min(max(backoff, header), self.max_delay).

Motivation and Context

No linked issue — found while auditing the retry path. The module's entire reason to exist is header-aware retry, and that path was dead.

Screenshots

Not applicable — no UI change.

API Changes

• This PR includes API changes

Change Type

• Bug fix
• New feature
• Performance improvement
• Refactoring
• Documentation
• Dependency/Build system
• Breaking change
• Other (specify):

Testing Performed

• Tested locally
• Manual/QA verification

Three tests were added to tests/unit/agents/new_chat/test_retry_after.py. Two build a real litellm.exceptions.RateLimitError rather than a hand-rolled fake, so they cover the shape production raises; the third pins the cap.

Run against the current dev code, before the fix:

FAILED tests/unit/agents/new_chat/test_retry_after.py::TestExtractRetryAfter::test_reads_headers_off_a_real_litellm_error
FAILED tests/unit/agents/new_chat/test_retry_after.py::TestExtractRetryAfter::test_reads_milliseconds_off_a_real_litellm_error
FAILED tests/unit/agents/new_chat/test_retry_after.py::TestDelayCalculation::test_caps_a_header_delay_at_max_delay
3 failed, 19 passed

After:

22 passed

No existing test changed. test_takes_max_of_backoff_and_header uses a 10s header under the default 60s cap, so the new clamp does not alter its result.

Full unit suite, before and after the change: 11 failed, 2997 passed, 1 skipped. The eleven are pre-existing on dev (git-tree and knowledge-store tests plus test_pat_fail_closed_static) and reproduce on a clean dev checkout on this machine — they are unrelated to this diff.

ruff check and ruff format --check are clean on both changed files.

What does not change

• Which exceptions are retried. _NON_RETRYABLE_CATEGORIES, _is_non_retryable and the retry_on hook are untouched.
• The exponential backoff formula, the jitter, the retry count, and the max_delay default of 60s.
• The surfsense.retrying custom event and its payload.
• Behaviour when a provider sends no header: _extract_retry_after_seconds still returns None and the backoff is used unchanged.

Remaining risk

Turns can now take longer, because a hint that was previously discarded is obeyed. That is the intended behaviour, and max_delay is its ceiling. I did not change the default cap — 60s is the existing value and picking a different one is a product call.

Checklist

• Follows project coding standards and conventions
• Documentation updated as needed
• Dependencies updated as needed
• No lint/build errors or new warnings
• All relevant tests are passing

High-level PR Summary

This PR fixes a bug in the retry middleware that prevented it from honoring provider Retry-After headers. The RetryAfterMiddleware was designed to respect rate limit hints from LLM providers but has never worked in production because it checked for dict headers when litellm actually returns httpx.Headers (a Mapping that is not a dict subclass). The fix changes the type check from isinstance(headers, dict) to isinstance(headers, Mapping) so the header-reading branch can execute. Additionally, the max_delay cap now applies to provider-supplied delays as well as exponential backoff to prevent extremely long sleep periods from blocking the chat turn, SSE stream, and database session.

⏱️ Estimated Review Time: 5-15 minutes

💡 Review Order Suggestion

Order	File Path
1	surfsense_backend/app/agents/chat/shared/middleware/retry_after.py
2	surfsense_backend/tests/unit/agents/new_chat/test_retry_after.py

[vercel]

@Yigtwxx is attempting to deploy a commit to the Rohan Verma's projects Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 40137550-f95b-461a-b0b4-3a0bb5cd6f14

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Yigtwxx]

Reading the red checks, since none of them come from this diff.

• Frontend Quality — FAILURE, and therefore Quality Gate. This PR touches no TypeScript at all. The failure is a single pre-existing format error in surfsense_web/app/(home)/free/[model_slug]/page.tsx, which is already on dev. The biome-check-web hook in .pre-commit-config.yaml sets always_run: true with pass_filenames: false, so it checks the whole surfsense_web tree no matter what a PR changed — the workflow's --from-ref/--to-ref narrowing does not reach it. Measured on a clean LF checkout of dev at a89216059: Checked 1097 files. Found 1 error. Every open PR inherits it.
• Journey — FAILURE. Fails in Build & start backend stack, before any test runs: container surfsense-e2e-celery_worker-1 is unhealthy. db, redis and backend all report healthy; only the worker times out. Infrastructure, not the diff.
• Vercel — FAILURE is Authorization required to deploy, the usual result for a fork PR, and recurseml/analysis — ERROR is the bot erroring on itself.

Green: Unit Tests, Integration Tests, Test Gate, Backend Quality, File Quality, Security Scan, CodeRabbit.

---

Update — both reds now have fixes rather than just explanations.

• The biome error is fixed in fix(web): use ES2022-safe array copies in client code #1707, whose Quality Gate came back green.
• The Journey bring-up is fixed in fix(e2e): probe the celery worker over the broker, not through the app #1708. Root cause: the worker healthcheck runs celery -A app.celery_app inspect ping, and -A makes the CLI import the entire application (embeddings, LiteLLM, docling) before it can send anything — about 40s in CI, against a timeout: 5s. Every probe was killed mid-import, so the worker was reported unhealthy while its own log showed celery@... ready. and no errors. That job has not passed in the last 60 runs, dev included.

Neither is a dependency of this PR; merging them first is just what turns this board green.