Skip to content

fix(api): cancel a sync retain when its client disconnects - #5358

Merged
nicoloboschi merged 2 commits into
mainfrom
fix/sync-retain-disconnect-cancellation
Oct 6, 2026
Merged

nicoloboschi merged 2 commits into
mainfrom
fix/sync-retain-disconnect-cancellation

Conversation

@nicoloboschi

Copy link
Copy Markdown
Collaborator

Fixes #4526.

The bug

ClientDisconnectCancellationMiddleware only monitored /memories/recall and /reflect, so
POST /memories never got a cancellation token. That one gate cost twice:

  • While queued — admit_for's dependency reads the token to let an abandoned request give up
    its place. For retain it read None, so a queued retain kept its lane slot after its caller
    had already hung up.
  • Once running — a sync retain (async: false) runs the whole extraction inline and creates
    no async_operations row, so _check_op_alive has nothing to key on either. Nothing could
    stop it.

Measured in the issue: a client that hung up at 3s left a retain running 1880s, holding its
LLM slot for the whole half hour and committing its document into a bank that had been deleted
85s earlier.

The fix

api/disconnect.py — _should_monitor takes the ASGI scope and also matches POST on
.../memories. DELETE shares that path and is not a retain, hence the method check;
/files/retain stays out because its body is streamed. The pump queue is bounded to one message
so monitoring a body-carrying route cannot hold a second copy of a retain body (base64
screenshots included) in memory; the token is tripped before the put, so a full queue never
delays cancellation.

api/http.py — new run_task_cancellable_on_disconnect, used by the sync retain branch.
The existing cooperative helper is the wrong shape here: recall and reflect spend their time in
worker threads that task cancellation cannot interrupt, which is why they check a token at stage
boundaries. A sync retain is the opposite — almost all of its wall time is one await, queued on
the LLM semaphore and then on the provider's response, so a checkpoint between stages never gets
a turn. This helper races the work as a task against the token, cancels it, awaits the unwind so
partial work still commits and subtasks die, then answers 499 and logs [RETAIN CANCELLED].
Cancelling a running retain is already a supported path: the worker runs every retain under an
absolute asyncio.timeout ceiling.

The admission half needed no code — with a token present, AdmissionAbandoned → 499 now fires
for a queued retain.

Tests

  • tests/test_recall_cancellation.py — _should_monitor gating including the DELETE and upload
    exclusions; the new helper stopping work that reaches no checkpoint while still running its
    finally; the teardown when the call is cancelled from above (shutdown), which asyncio.wait
    does not propagate on its own; an error propagating through; and an end-to-end ASGI test
    driving the real app, where the client hangs up and the route answers 499 with the extraction
    unwound.
  • hindsight-system-tests/tests/test_18_retain_disconnect.py — the blackbox story, through the
    published client: a client with a short timeout walks away while the extraction is held open
    (llm.hold), and nothing is committed. Verified red before the fix (documents.total == 1,
    the exact retain with async:false is never cancelled on client disconnect — an abandoned request ran 1880 s and wrote into a bank deleted while it ran #4526 symptom) and green after. Plus the companion assertion that a patient caller
    still gets its whole retain, so the guard can't pass by cancelling everything.

No OpenAPI drift (generate-openapi.sh is a no-op); the control-plane retain proxy already
passes the dataplane status through.

Deliberately not in scope

Both are separate root causes from the one in the title:

  • A long sync retain can still outlive a bank delete. It needs a bank-liveness check, which
    _check_op_alive cannot provide here (no operation row to key on), and a single-chunk retain
    has no sub-batch boundary to check at. The disconnect fix covers the case the issue measured;
    it does not cover a caller who waits patiently through a concurrent delete.
  • Sync retain still has no ack or idempotency key, so a caller that times out cannot tell
    "still running" from "lost" and cannot retry safely. The issue raises refusing async: false
    outright as an alternative — that would be a breaking change to the documented default, so it
    belongs in its own discussion.

No figure: the change is a guard on one request path, and the issue's own timeline already shows
the flow better than a diagram would.

`ClientDisconnectCancellationMiddleware` only monitored `/memories/recall`
and `/reflect`, so `POST /memories` never got a cancellation token. That one
gate cost twice: the admission dependency saw `abandoned=None`, so a queued
retain kept its place after its caller left, and nothing could stop the
extraction once it started — a sync retain (`async: false`) has no
`async_operations` row, so `_check_op_alive` has nothing to key on either.
Measured in #4526: a client that hung up at 3s left a retain running 1880s,
holding its LLM slot and committing into a bank deleted while it ran.

- `_should_monitor` now takes the ASGI scope and also matches `POST` on
  `.../memories`. DELETE shares that path and is not a retain, hence the
  method check; `/files/retain` stays out, its body is streamed.
- The pump queue is bounded to one message, so monitoring a body-carrying
  route cannot hold a second copy of a retain body in memory. The token is
  tripped before the put, so a full queue never delays cancellation.
- New `run_task_cancellable_on_disconnect`, used by the sync retain branch.
  The cooperative helper is the wrong shape here: a sync retain spends its
  wall time on a single await — queued on the LLM semaphore, then on the
  provider's response — so a checkpoint between stages never gets a turn.
  This one races the work as a task against the token, cancels it, awaits
  the unwind so partial work still commits and subtasks die, then answers
  499. Cancelling a running retain is already a supported path: the worker
  runs every retain under an absolute `asyncio.timeout`.

The admission half needed no code — with a token present, `admit_for`'s
`AdmissionAbandoned` -> 499 now fires for a queued retain.

Tests: `_should_monitor` gating including the DELETE and upload exclusions;
the new helper stopping work that reaches no checkpoint and still running
its `finally`; and an end-to-end ASGI test driving the real app, where the
client hangs up and the route answers 499 with the extraction unwound.

Not addressed, and separate root causes from this one: a long sync retain
can still outlive a bank delete (no operation row to key liveness on, and a
single-chunk retain has no boundary to check at), and sync retain still has
no ack or idempotency key.

Closes #4526
- `run_task_cancellable_on_disconnect` cancels the work when the call is
  itself cancelled from above (shutdown, a deadline). `asyncio.wait` does
  not propagate its own cancellation to what it waits on, so the retain
  would have kept running with nobody left to answer it.
- Updated the admission dependency's comment and `disconnect.py`'s module
  docstring, both of which still said recall and reflect were the only
  routes carrying a token.
- Added the blackbox story `test_18_retain_disconnect.py`: a client with a
  short timeout walks away while the extraction is held open, and nothing
  is committed. Verified red before the fix (document total 1 != 0) and
  green after. Plus the companion assertion that a patient caller still
  gets its whole retain.
- Tests for the teardown and for an error propagating through the helper;
  `_` for the unused unit-ids result; clearer test name.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

retain with async:false is never cancelled on client disconnect — an abandoned request ran 1880 s and wrote into a bank deleted while it ran

1 participant