[PR #99] [MERGED] feat: collect a Router deadline 504 under the same Idempotency-Key by default #99

New issue

Closed

opened 2026-10-31 22:01:22 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-31 22:01:22 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-python-sdk/pull/99
Author: @mattmillerai
Created: 8/27/2026
Status: ✅ Merged
Merged: 8/27/2026
Merged by: @mattmillerai

Base: main ← Head: matt/be-9836-router-deadline-504-same-key-retry

📝 Commits (2)

d0f22dc feat: collect a Router deadline 504 under the same Idempotency-Key by default
8d0d883 fix: close the gates the collect loop is only safe behind

📊 Changes

10 files changed (+1078 additions, -153 deletions)

View changed files

📝 README.md (+50 -25)
📝 src/comfy_low/errors.py (+43 -1)
📝 src/comfy_low/transport.py (+10 -1)
📝 src/comfy_sdk/models.py (+31 -11)
📝 src/comfy_sdk/retry.py (+296 -69)
📝 src/comfy_sdk/router_exceptions.py (+18 -13)
📝 tests/conftest.py (+63 -14)
📝 tests/test_error_mapping.py (+79 -2)
📝 tests/test_models_run_retry.py (+479 -15)
📝 tests/test_router_exceptions.py (+9 -2)

📄 Description

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.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-python-sdk/pull/99 **Author:** [@mattmillerai](https://github.com/mattmillerai) **Created:** 8/27/2026 **Status:** ✅ Merged **Merged:** 8/27/2026 **Merged by:** [@mattmillerai](https://github.com/mattmillerai) **Base:** `main` ← **Head:** `matt/be-9836-router-deadline-504-same-key-retry` --- ### 📝 Commits (2) - [`d0f22dc`](https://github.com/Comfy-Org/comfy-python-sdk/commit/d0f22dc80688bdc91321c3c0efc71efbb6c3195e) feat: collect a Router deadline 504 under the same Idempotency-Key by default - [`8d0d883`](https://github.com/Comfy-Org/comfy-python-sdk/commit/8d0d883d47ad03a3404643eb622e5583a1ee47d2) fix: close the gates the collect loop is only safe behind ### 📊 Changes **10 files changed** (+1078 additions, -153 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+50 -25) 📝 `src/comfy_low/errors.py` (+43 -1) 📝 `src/comfy_low/transport.py` (+10 -1) 📝 `src/comfy_sdk/models.py` (+31 -11) 📝 `src/comfy_sdk/retry.py` (+296 -69) 📝 `src/comfy_sdk/router_exceptions.py` (+18 -13) 📝 `tests/conftest.py` (+63 -14) 📝 `tests/test_error_mapping.py` (+79 -2) 📝 `tests/test_models_run_retry.py` (+479 -15) 📝 `tests/test_router_exceptions.py` (+9 -2) </details> ### 📄 Description ## 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. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-31 22:01:22 +00:00