fix(gateway): tombstone deleted sessions to stop trace-store leak

[jeffreysijuntan]

Problem

The gateway's in-memory trace store (--store memory) leaks sessions during async RL training. Observed live on a Fireworks run: ~157 orphaned sessions/hr, each exactly 1 trace, at old weight versions, with large mid-episode prompts — accumulating unbounded (the store has no TTL/eviction).

Root cause — a persist-after-delete race

Every rollout correctly deletes its own gateway session in a finally (so filtering, staleness, and errors are all cleaned up — those were not the leak). The leak is a lifecycle race:

1. Traces persist fire-and-forget (sync_traces=False): the proxy does asyncio.create_task(store_trace(...)) after each completion returns.
2. When a rollout ends it calls flush() then delete_session(uid). flush() only drains persist tasks that already exist — an LLM call still in flight at the provider hasn't reached the persist line yet, so it can't be flushed.
3. That straggler returns after the delete → store_trace runs → and MemoryTraceStore resurrects the session because self._session_index[session_id] is a defaultdict(list) with no tombstone. Nothing deletes that uid again → permanent orphan with exactly 1 trace.

Signature match: single trace (the lone straggler), large prompt (slow _build_trace_data widens the window under high concurrency), grouped by task-uuid (a hard prompt makes several group members straggle together), early weight version (early policy aborts right after big tool-call turns).

Fix — a permanent tombstone

Session ids are unique uuids (uid = f"{uuid4()}:{rollout_idx}") that are never reused, so a deleted session must be rejected permanently, not for a time window:

• delete_session records the session_id in a tombstone set.
• store_trace drops any write for a tombstoned session instead of resurrecting it — always correct, since a deleted uid is never legitimately written again.
• The tombstone set is FIFO-capped (max_tombstones, default 1,000,000) purely as a memory backstop for long-lived multi-run gateways. The cap is far larger than the number of sessions deletable while any request is still in flight, so eviction never exposes a live straggler. max_tombstones=0 disables the guard (legacy behaviour).

Why not require pre-creation instead? Sessions are implicitly created on first request (a documented gateway feature for captured/external sessions), and the persist path (ReverseProxy) doesn't hold the SessionManager's liveness state — so remembering deletions (a tombstone) is the correct store-level mechanism.

Why not a TTL? A time window has a real failure mode — if a request outlives it (the heartbeat budget allows up to 3600s), a very-late straggler could still resurrect the session. A permanent block has no such edge.

Cost: one interned session_id per deleted session (~tens of MB worst-case for an entire multi-day run, hard-capped); no change to the read/serialization path.

Tests

tests/unit/test_store.py — new TestTombstone: straggler-after-delete is dropped (not resurrected), the block is permanent across repeated late writes, unrelated sessions unaffected, FIFO cap evicts the oldest tombstone, and max_tombstones=0 restores legacy behaviour. Full store suite: 37 passed.

Scope / follow-up

SqliteTraceStore shares the same resurrection pattern (INSERT OR REPLACE/IGNORE on write, plain DELETE) and would take the identical guard — left as a follow-up since the observed leak is on the in-use memory backend.

This is a separate fix from the array-packing perf change (#787, already merged into this branch); the leak is a distinct, slower, unbounded issue on top of the O(N²) working-set growth.

🤖 Generated with Claude Code