Cancel scheduler session on HTTP disconnect / drain error

[FluffyAIcode]

Summary

Server-side fix for the orphan-session bug surfaced by the 2026-05-30
Mac M4 4-hour bench_long_session.py run. This is the second of the
two PRs needed to unblock a clean re-run; PR #24 fixed the bench-side
metric-name mismatch and added the scheduler_kv_live_bytes gauge.

The aborted run report (results/platform-tests/bench_long_session_mac_1780130542.aborted.json
on AgentMemory/bench-long-session-mac-results-8e7f) recorded:

• 58 successful turns over 29.4 minutes
• Server-side logs: 3+ hours of POST /v1/chat/completions 429 Too Many Requests
• Bench client showed 0 errors in errors[] because it was hung on a
  single in-flight HTTP request

Diagnosis: when a non-streaming chat request was cancelled or timed
out client-side, the scheduler worker kept running until generation
finished naturally, holding the only slab. With --max-concurrent 1
every subsequent request returned 429.

Changes

inference_engine/scheduler/scheduler.py (+5 / -1 lines)

The cross-thread token enqueue inside _run_session now blocks on
the run_coroutine_threadsafe future before returning. Without this,
the worker thread can keep producing tokens while the event loop
serially drains them, but on generate() return the terminal sentinel
is pushed back on the loop and can race ahead of still-pending token
puts — corrupting the SSE chunk order.

enqueue = asyncio.run_coroutine_threadsafe(
    session.token_queue.put(int(tok_id)), loop
)
enqueue.result()  # ← preserve token-before-sentinel ordering

inference_engine/server/app.py (+38 / -5 lines)

Two changes to the non-streaming /v1/chat/completions path:

1. Cancel on disconnect / cancellation. A new helper
   _collect_non_streaming_tokens polls
   request.is_disconnected() every 50 ms while draining the
   session's token queue, and calls scheduler.cancel_session(session)
   when the client has gone away. Streaming responses already had
   this loop; the non-streaming path was the orphan-session source.
2. Cancel on engine error. The BaseException catch around the
   drain now calls scheduler.cancel_session(session) before
   converting to HTTP 500, so a verifier crash mid-generate also
   releases the slab.
3. The asyncio.CancelledError path (transport-level cancel from
   the test client / a real timed-out client) cancels the session
   and re-raises, so the route handler never returns from a
   half-drained generation with the slab still held.

tests/inference_engine/server/test_app_streaming.py (+32 lines)

test_collect_non_streaming_tokens_cancels_on_disconnect exercises
the new helper directly with a _FakeRequest that returns
is_disconnected() == True on first poll; asserts the session
state transitions to CANCELLED and scheduler.active_count
returns to zero. Real concrete classes only — no mocks.

Verification

$ pytest tests/inference_engine/server/ tests/inference_engine/scheduler/ -q
278 passed in 3.98s

$ pytest tests/inference_engine/ -q   (with PR #24 merged on top)
390 passed in 10.57s

Compatibility with PR #24

Both PRs touch inference_engine/server/app.py but on different
lines (PR #24: bootstrap snapshot + /metrics scrape; this PR: non-
streaming chat handler). I tested the auto-merge locally:

$ git merge origin/AgentMemory/scheduler-orphan-session-fix-8e7f \
       (onto AgentMemory/bench-long-session-fixes-8e7f)
Auto-merging inference_engine/server/app.py
Automatic merge went well

$ pytest tests/inference_engine/ -q
390 passed

Merge order does not matter; either PR can land first.

Recommended re-run path

After both PR #24 and this PR land, the validation sequence is:

# 30-min short run first
PYTHONPATH=. python3 scripts/bench_agentic/bench_long_session.py \
    --kakeya-url http://127.0.0.1:8000 --kakeya-model kakeya-v1 \
    --duration-s 1800 --turn-spacing-s 5 --max-tokens 64 \
    --report results/platform-tests/bench_long_session_mac_short_$(date +%s).json

Success criteria:

• All turns record non-empty metrics dict with 5 keys including
  scheduler_kv_live_bytes. (PR Fix bench_long_session metric names + expose scheduler_kv_live_bytes gauge #24)
• scheduler_kv_live_bytes stays bounded — verifies ADR 0006 §2.3.
• No sustained 429. This PR is the gate on that one.
• agg.kv_bounded == True in the final JSON.

If the short run is clean, a 4-hour run is the GA-evidence step.

References

• Triage data: results/platform-tests/bench_long_session_mac_1780130542.aborted.json
  on AgentMemory/bench-long-session-mac-results-8e7f
• Companion PR: Fix bench_long_session metric names + expose scheduler_kv_live_bytes gauge #24 (bench metric names + scheduler_kv_live_bytes gauge)
• ADR 0003 — slab pool / verifier integration
• ADR 0006 §2.3 — long-session memory-stability claim