fix(inference): cancel flex requests when the caller disconnects

[sejori]

Problem

Blocking and streaming flex requests (service_tier: "flex" on /v1/chat/completions, or non-background /v1/responses) are written into fusillade and polled by the handler. Nothing cancelled the row when the client dropped the connection. The daemon claimed and dispatched the abandoned request anyway, and because upstream timeouts are retriable it could re-dispatch it repeatedly. A client with its own timeout that retries therefore stacks redundant work onto an already backlogged engine.

The daemon could not have honoured a cancel even if one had been written: only batched requests have a cancellation path (batch cancelling_at → shared token). Batchless requests got a fresh token nobody fired (daemon/mod.rs, None => CancellationToken::new()).

Change

dwctl

• flex_stream_response selects on the SSE channel closing (the receiver lives in the response body, dropped on disconnect) and cancels the row.
• The blocking handlers (handle_flex foreground, handle_chat_completion_flex) hold an AbandonGuard across the poll. hyper drops the handler future on client disconnect, so the guard's Drop issues the cancel. It is disarmed once a terminal state has been rendered.
• Both also cancel when the poll itself gives up (1h timeout / storage error), since the client has by then been answered with an error.

fusillade daemon

• Per-request cancellation tokens for in-flight batchless requests, removed when the processing task exits.
• The existing cancellation poll additionally asks storage which of those rows are now canceled and fires their tokens, which aborts the HTTP task through the existing complete() select. New gauge fusillade_cancellation_poll_requests_checked.

storage

• Storage::cancel_batchless_request(id) -> bool: a single guarded UPDATE mirroring the persist(Canceled) transition (only completed/failed resist). Needed because the default cancel_requests goes through get_requests, which inner-joins batches and never sees batchless rows.
• Storage::get_cancelled_request_ids(ids): batchless sibling of get_cancelled_batch_ids.

canceled remains the soft terminal: a completion that lands anyway supersedes it and is billed, exactly as for batches today.

Behaviour change

A foreground flex request whose caller disconnects is now cancelled rather than left to complete. Its result can no longer be fetched later by ID. Background (background: true) requests are unaffected: they return 202 immediately and are never polled by the handler.

Tests

• fusillade-arsenal: get_cancelled_request_ids_reports_only_canceled_rows covers both new storage methods, including that a completed row resists the cancel.
• fusillade integration: cancelling_in_flight_batchless_request_aborts_http runs the daemon against a mock upstream that never answers, cancels the row, and asserts the in-flight call is aborted, not retried, and the row stays canceled.
• dwctl: dropping_the_stream_cancels_the_flex_request, abandon_guard_cancels_when_handler_is_dropped, abandon_guard_leaves_delivered_result_alone.

Not covered here

Aborting the daemon's HTTP task closes the loopback hop, which drops the upstream call from onwards. Whether the inference engine actually stops generating on a dropped connection is the same open question as for realtime disconnects and is outside this PR.

🤖 Generated with Claude Code

[cloudflare-workers-and-pages]

Deploying control-layer with    Cloudflare Pages

Latest commit:	d90fc22
Status:	✅  Deploy successful!
Preview URL:	https://7ade9635.control-layer.pages.dev
Branch Preview URL:	https://fix-flex-cancel-on-disconnec.control-layer.pages.dev

View logs

[Copilot]

Copilot review overview

🟡 Changes recommended

Enqueue-time disconnect races and a potentially stale-replica cancellation check can leave abandoned requests running.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 3 · 1

Open (4)

• Arm cancellation guard before enqueue · New
• Cover enqueue phase with cancellation guard · New
• Use primary executor for cancellation detection · New
• Update polling gauges before empty-list continue · New

What changed in this PR

Adds cancellation for abandoned foreground flex inference requests across dwctl, Fusillade, and storage.

Changes:

• Cancels requests on client disconnects and handler timeouts.
• Adds batchless cancellation tokens and daemon polling.
• Adds storage APIs and cancellation tests.

File	Summary
fusillade/​tests/​integration.rs	Tests cancellation of in-flight batchless requests.
fusillade/​src/​daemon/​mod.rs	Tracks and triggers batchless request cancellation.
fusillade-core/​src/​manager.rs	Extends the storage cancellation interface.
fusillade-arsenal/​src/​postgres.rs	Implements batchless cancellation queries.
dwctl/​src/​inference/​streaming.rs	Cancels streaming requests on disconnect.
dwctl/​src/​inference/​store.rs	Adds abandonment cancellation helpers and guards.
dwctl/​src/​inference/​middleware.rs	Applies guards to blocking flex handlers.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[cubic-dev-ai]

1 issue found across 7 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="fusillade-arsenal/src/postgres.rs">

<violation number="1" location="fusillade-arsenal/src/postgres.rs:5197">
P1: A retriable upstream result can overwrite this cancellation and re-enter the retry path, because batchless rows have no batch `cancelling_at` fence. Preserve a durable cancellation fence through retry decisions or suppress retries for rows canceled by `cancel_batchless_request`; otherwise a disconnected flex request can be dispatched again.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P1: A retriable upstream result can overwrite this cancellation and re-enter the retry path, because batchless rows have no batch cancelling_at fence. Preserve a durable cancellation fence through retry decisions or suppress retries for rows canceled by cancel_batchless_request; otherwise a disconnected flex request can be dispatched again.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At fusillade-arsenal/src/postgres.rs, line 5197:

<comment>A retriable upstream result can overwrite this cancellation and re-enter the retry path, because batchless rows have no batch `cancelling_at` fence. Preserve a durable cancellation fence through retry decisions or suppress retries for rows canceled by `cancel_batchless_request`; otherwise a disconnected flex request can be dispatched again.</comment>

<file context>
@@ -5156,6 +5156,56 @@ impl<P: PoolProvider> Storage for PostgresRequestManager<P> {
+                canceled_at = NOW()
+            WHERE id = $1
+              AND batch_id IS NULL
+              AND state NOT IN ('completed', 'failed', 'canceled')
+            RETURNING id
+            "#,
</file context>

[sejori]

Not reachable. A retriable failure on a cancelled row goes through reschedule_for_retry, whose UPDATE is fenced on state IN ('claimed', 'processing') AND daemon_id = $2 AND claimed_at = $5 (postgres.rs reschedule_for_retry). A row moved to canceled by cancel_batchless_request fails that fence, the daemon logs the retry as not persisted, and the row stays canceled. When retries are exhausted, persist(Failed) does supersede canceled — that is the existing soft-terminal rule for batches as well, and it ends the work rather than re-dispatching it. The integration test cancelling_in_flight_batchless_request_aborts_http asserts the mock upstream is called exactly once after the cancel.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Cancellation can race with request creation, potentially leaving abandoned work pending.

Review effort: Lite
Findings: None

Resolved since last review (4)

• Use primary executor for cancellation detection
• Cover enqueue phase with cancellation guard
• Arm cancellation guard before enqueue
• Update polling gauges before empty-list continue

Previously missed (1)

In code that hasn't changed since last review

Record detached cancellation failures in background error metrics

dwctl/​src/​inference/​store.rs:325

This failure is handled by a detached task after the request future is dropped, so the cancellation error is not reflected in the HTTP metrics and can leave abandoned work running. The repository's background-task convention is to use background_error! for swallowed off-request-path failures (see dwctl/src/metrics/errors.rs:3-9); please record this with an appropriate component/reason label so failed cancellations are observable and alertable.

[cubic-dev-ai]

All reported issues were addressed across 5 files (changes from recent commits).

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[sejori]

Re Copilot's "previously missed" note on cancel_abandoned_request (store.rs, not attached to a thread): done in 15c2458. The Err arm now goes through background_error! with a new flex_cancel component and reason cancel_abandoned, severity Error, so failed cancels land in dwctl_background_errors_total.

[cubic-dev-ai]

1 issue found across 1 file (changes from recent commits).

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="dwctl/src/test/responses.rs">

<violation number="1" location="dwctl/src/test/responses.rs:1405">
P2: `drop(response)` runs after the complete 202 exchange, and the immediate mock lets the daemon finish before the detached cancellation can be observed. Hold the upstream request in flight and assert the row remains non-canceled before allowing completion; otherwise a background-cancellation regression can pass this test.</violation>
</file>

<sub>Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: drop(response) runs after the complete 202 exchange, and the immediate mock lets the daemon finish before the detached cancellation can be observed. Hold the upstream request in flight and assert the row remains non-canceled before allowing completion; otherwise a background-cancellation regression can pass this test.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At dwctl/src/test/responses.rs, line 1405:

<comment>`drop(response)` runs after the complete 202 exchange, and the immediate mock lets the daemon finish before the detached cancellation can be observed. Hold the upstream request in flight and assert the row remains non-canceled before allowing completion; otherwise a background-cancellation regression can pass this test.</comment>

<file context>
@@ -1373,3 +1373,48 @@ async fn test_realtime_zdr_suppresses_stored_bodies(pool: PgPool) {
+        .expect("202 body carries a resp_<uuid> id");
+
+    // The client has its id and moves on: the HTTP exchange is over.
+    drop(response);
+    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
+
</file context>

[sejori]

The premise doesn't hold here: the test harness sets batch_daemon.enabled = DaemonEnabled::Never, so no daemon ever claims the row and the mock is never reached. The row can only leave pending if something cancels it. Tightened in d90fc22 to assert_eq!(state, "pending") so a spurious cancel fails the test outright, and added a comment saying why.