feat: collect a Router deadline 504 under the same Idempotency-Key by default

[mattmillerai]

ELI-5

When a model run takes longer than Comfy is willing to hold the connection open, the server answers 504 deadline_exceeded — but it keeps your idempotency reservation, stamps the provider's task handle on it, and tells you (in Retry-After) when to ask again with the same Idempotency-Key. Asking again collects the generation that is still running rather than starting a second one. The SDK was not asking. It filed that 504 with every other 5xx as "outcome unknown, key possibly claimed" and handed the caller an error for a generation they get billed for either way. Now one run() call waits out the pace, re-presents the same key, and rides the collect loop to the finished result.

What changed

A fourth failure class — collectable — sits ahead of the unknown-outcome one in RetryPolicy.should_retry:

• a deadline_exceeded 504 carrying Retry-After, and
• a generation_in_progress 409 carrying Retry-After — the idempotency layer's "still in progress", which is what the collect retry meets when it arrives before the generation finishes.

Both gates are load-bearing, and both come straight off the vendored contract in spec/router-openapi.yaml:

• the bucket, not the status. deadline_exceeded shares 504 with provider_timeout, and a header-less 504 from an intermediary is read as the latter — where nothing blesses a same-key resend. Only the bucket the contract names is resent. (This is the exact reason the previous PR gave for not acting on the bucket; error_bucket_of is what removes it.)
• the pace. The spec says Retry-After is present on this bucket "only when Comfy holds a handle to a generation the provider is still running… Absent when there is nothing to collect". No pace, nothing to collect, no resend.

error_bucket_of reads the bucket off either layer's name for it — error_type on a typed RouterError, code on the comfy_low.ApiError that POST /models/run raises today. Reading only error_type would have made the whole rule unreachable on the route it was written for: a retry that is a silent no-op with no test failing, which is the failure mode this module has already been bitten by twice.

retry_possibly_in_flight is untouched and keeps its meaning — the class no contract characterises (a 5xx that named no collectable pace, and the client-side timeout). The new class is the one where the server itself said the resend is safe. Both docstrings now say which is which.

Two budgets, because the classes have two shapes. max_elapsed stays at 60s and governs the fast classes (connect failure, paced 429), which resolve in seconds or not at all — a blackholed host has no business pinning a caller thread for longer. A new collect_max_elapsed is 1200s and governs the collect class alone, selected per failure by is_collectable; both are measured from the same Retrier origin.

1200s is derived, not invented: a deadline_exceeded 504 arrives at Comfy's own bound, the same ten minutes as MODEL_RUN_TIMEOUT, so a budget of one deadline window is already spent by the time the 504 lands and the collect attempt it exists for never starts. Two windows is one to reach the 504 and one to collect what it left running. Nothing else pays for that room.

No feature flag, as asked. The control surface is the existing RetryPolicy: Comfy(retry=NO_RETRY) disables retrying entirely, and RetryPolicy(retry_collectable=False) turns off just this rule. The new field is declared last in the dataclass, deliberately out of reading order, so that adding it cannot silently change what an existing positional RetryPolicy(...) means.

Docstring correction. models.run() claimed the key is "single-use and reject-on-duplicate". That is the v2 jobs API's rule (spec/openapi.yaml's shared IdempotencyKey parameter), and it governs submit(), not this route. It is scoped there now, and the router's replay-on-duplicate behaviour — same key + same body collects the generation already started, different body is rejected 422 idempotency_key_reuse — is stated in its place. The README's retry table and budget row were carrying the same claim and are corrected too.

Judgment calls

• The 409 rule is not scoped to "only after a 504." Making it stateful (retry a 409 only if the previous attempt was a deadline 504) would buy very little and cost real complexity in the retry loop. It is gated on the generation_in_progress bucket and on Retry-After instead, the same two-gate shape the 504 uses, and it fails closed: a bucket-less proxy or WAF conflict, hash_mismatch (which spec/openapi.yaml gives a Retry-After and which is still a deterministic refusal), and asset_in_use all stay refusals. A 409 is a 4xx, so this is the one place the diff narrows an existing "every other 4xx is never retried" branch; the narrowing is bucket-scoped, and it is tested in both directions.
• The accepted risk is deployment skew. Against a Router deployment whose reservation does not survive the deadline 504, the collect retry comes back 422 idempotency_key_reuse and the caller sees that instead of the 504 — one extra request, and a less diagnosable error. That is the same objection the previous PR raised against retrying service_unavailable by default, and the thing that answers it here is that the contract states the carry for this bucket only. retry_collectable=False is the escape hatch, and it is named in the DeadlineExceeded docstring where someone hitting the 422 will look.
• service_unavailable is deliberately still not retried by default. Nothing in either vendored spec says a 503 releases the key; that decision, and its test, are unchanged.
• No capability is denied by this diff (the falsification trigger): it turns a retry on. The two "not retried" assertions it adds — a 504 without a pace, and a provider_timeout 504 — are the contract's own distinctions, read out of spec/router-openapi.yaml in this repo rather than assumed.

Scope not covered by the fix

Three SDK surfaces send an Idempotency-Key: POST /models/run (this change), POST /jobs behind submit(), and POST /assets. The other two are the v2 jobs contract's reject-on-duplicate surface with no documented replay, and they keep their own 429 handling — untouched, and correctly so. Of the fifteen router error buckets the vendored spec declares, exactly one (deadline_exceeded) carries a contract-level same-key-retry blessing and is the one this enables; a second (service_unavailable) asks for a retry but says nothing about the key and stays behind the opt-in; the remaining thirteen are unchanged.

Testing

pytest -q — 574 passed, 4 skipped (the same 4 gateway-e2e skips as on main). ruff check ., ruff format --check ., mypy src, scripts/check_public_repo_hygiene.py and scripts/check_drift.py (both model and router-error-type checks) all clean.

Collect-loop tests: a deadline 504 + Retry-After is collected under one key, sync and async, with both attempts asserted to carry the same non-None header value; a 504 with no pace is raised, not retried; a provider_timeout 504 with a pace is raised, not retried; a generation_in_progress 409 + Retry-After is waited out across three attempts on one key; a 409 with no pace, and a paced hash_mismatch 409, are both still refusals; NO_RETRY and retry_collectable=False each disable the loop (the latter now also with retry_possibly_in_flight=True, where it previously did not hold); the loop gives up at collect_max_elapsed (asserted both on the wire and exactly against a fake clock, where a 30s pace is clamped to the 10s left of a 100s budget); and the bucket is read off both error_type and code.

Gate tests added in review: a Retry-After: 0 falls back to the jittered backoff instead of spinning, while still counting as the server holding a handle; the collect budget outlasts one deadline window and applies only to the collect class; a 504 arriving at the 600s bound still gets a collect attempt; the body snapshot survives a caller mutating a nested value between attempts (verified to fail on the shallow copy); and error_from_envelope reads Router's {detail, error_type} + X-Comfy-Error-Type shape without retyping any status this API already has a documented code for.

The stub server gained model_run_collects_after_deadline — the narrower carry the router describes, where only the 504's reservation survives — kept distinct from model_run_replays_idempotency_key, the whole-deployment replay the opt-in exists for. That carve-out is keyed on the bucket as well as the status, so the stub is not more permissive than the contract it stands in for. It also gained model_run_router_error_shape, which answers in Router's own error shape rather than the v2 envelope; without it, nothing exercised the shape a real deployment actually sends. Without the carry flag the stub still claims the key on a 5xx and answers the retry 422, and test_the_default_collect_loop_against_a_non_collecting_deployment pins exactly that outcome: the assertion is the contract, not a stub loosened to fit the change.

Residual

Two items previously listed here — a server-named pace of 0 spinning for the whole budget, and the larger budget being charged to every failure class — were fixed in review (8d0d883) rather than carried, along with six other panel findings: the bucket never reaching is_collectable against Router's real error shape (which made the whole rule a production no-op), the ungated 409, the shallow body snapshot, the ineffective retry_collectable=False under the in-flight opt-in, the stub's status-only key-claim carve-out, and the README table row.

• The Router idempotency middleware's own source was not read. It lives in the backend repo, which is not reachable from where this was written, so the parts of the behaviour that come from it rather than from spec/router-openapi.yaml are taken on report: the 409 "already in progress" + Retry-After response shape, the 24h replay window, and same-body-collects / different-body-422. The 504 half is fully grounded — the vendored spec states the bucket's same-key retry and its Retry-After semantics verbatim — but the 409 rule and the corrected models.run() docstring rest on the unread half. Anyone with access should confirm the 409 status and header against the middleware before this is treated as settled.
• The vendored Router spec does not declare any of it. spec/router-openapi.yaml's runRouterModel operation declares no Idempotency-Key parameter and no 409 response, so the idempotency layer's contract is invisible to the SDK's own drift gate. The next Router spec sync should add both; until then tests/test_models_run_retry.py is the only thing pinning the 409 behaviour, and it pins this PR's premise rather than the server's.
• A collect retry that hits 422 idempotency_key_reuse replaces the original 504. The retry loop raises whatever the last attempt raised, so on a deployment without the reservation carry the caller loses the diagnosable error. test_the_default_collect_loop_against_a_non_collecting_deployment now pins that outcome rather than leaving it undescribed, and retry_collectable=False is the escape hatch. Preserving and re-raising the first failure when the retry fails with a key-reuse error would fix it properly, at the cost of state in the loop; filed as a follow-up rather than attempted here, since it changes what the loop raises for every class, not just this one.
• The TypeScript SDK twin was not touched or checked. It is a separate repo and out of scope for this change; the cross-SDK parity check will flag the drift this introduces until it lands there.

Provenance

• Authored by: agent-work loop
• Verified: pytest -q: 574 passed, 4 skipped; ruff check .: All checks passed; ruff format --check .: 51 files already formatted; mypy src: no issues in 19 files; scripts/check_drift.py: models in sync with spec/openapi.yaml, all 15 router error types covered; scripts/check_public_repo_hygiene.py: no internal-only references
• Deviations: none of the acceptance criteria were skipped. The two that are only partly grounded (the 409 rule and the replay-semantics docstring) are named under Residual with what was not read; the 409 rule is now bucket-gated and fails closed, which narrows what rests on the unread half. The default max_elapsed change stated in the original acceptance criteria (60s → 600s) was superseded in review: max_elapsed stays at 60s and the collect class gets its own collect_max_elapsed, which meets the criterion's intent (a collect loop that outlasts the deadline) without charging the other classes for it.

[coderabbitai]

Important

Approval pending

CodeRabbit has no unresolved comments, but it has not reviewed the latest commit.

Use the checkbox below to review the latest commit. CodeRabbit will approve the changes if it finds no blocking issues.

• 🔍 Trigger review

---

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

[github-actions]

🔍 Cursor Review — Consolidated panel

Triggered by @mattmillerai.

Found 10 finding(s).

Severity	Count
🟠 High	3
🟡 Medium	4
🟢 Low	3

Panel: 8/8 reviewers contributed findings.

[robinjhuang]

Approved — full autonomy check passed.

---

Generated by Claude Code

[mattmillerai]

All ten open review threads addressed in 8d0d883 and resolved with per-thread replies. Full suite green: 574 passed, 4 skipped; ruff check, ruff format --check, mypy src, check_drift.py and check_public_repo_hygiene.py all clean.

The two that changed the shape of the PR:

The collect rule was a silent no-op against a real Router 504. POST /api/v2/models/run is fronted by Router, whose error body is {detail, error_type} with the bucket on X-Comfy-Error-Type — not the v2 {error: {code}} envelope error_from_envelope read. Every real 504 reached is_collectable with the status-derived "error" bucket and fell through, while all eight existing tests passed because the stub only ever spoke the envelope. I reproduced it by reverting just the fallback: the new wire test fails with the raw 504. error_from_envelope now reads the header and body error_type, placed after _CODE_BY_STATUS so a Router 429 still maps to queue_full/QueueFull and nothing integrators already catch gets retyped — only the 5xx range, where the status determines nothing, is left to the bucket.

The budget is split rather than raised. Two findings pulled opposite ways — one that 600s could not outlast the 600s deadline window it was sized for, one that a policy-wide 10x hurt every fast class. Both are right, so max_elapsed goes back to 60s for the fast classes and a new collect_max_elapsed (1200s, two deadline windows) governs the collect class alone, selected per failure by is_collectable.

The rest: the 409 branch is bucket-gated on generation_in_progress and fails closed (a paced hash_mismatch stays a refusal); Retry-After: 0 falls back to the jittered backoff instead of spinning, while the header's presence still counts as the server holding a handle; retry_collectable=False now holds under retry_possibly_in_flight=True; the body snapshot is deep; the stub's key-claim carve-out takes the bucket as well as the status; and the README table row moved. Coverage added for the non-collecting deployment the default bets against.

One item deferred rather than fixed here: when a same-key retry is rejected 422 idempotency_key_reuse, that 422 replaces the original failure. Real and reachable, but it changes what the retry loop raises for every class, so it wants its own diff — filed as a follow-up, and the outcome is now pinned by a test rather than left undescribed.

Not merging — leaving that to a human reviewer.

[mattmillerai]

🤖 The reviews loop filed Linear follow-up ticket(s) for review thread(s) deferred as out of scope for this PR:

• BE-9948 — Preserve the original failure when a same-key retry is rejected with idempotency_key_reuse — filed as agent-spike (premise unverified)

The following carry agent-spike instead of agent-ok because their reachability claim was not backed by evidence (BE-5378) — the claim is investigated before any code is written, and "the premise does not hold" is a valid, successful outcome:

• Preserve the original failure when a same-key retry is rejected with idempotency_key_reuse — no reachability block in the proposal

[robinjhuang]

Approved — full autonomy check passed.

---

Generated by Claude Code