From d27eb59e13bfb91c8588a5a3d141edc7828c7b21 Mon Sep 17 00:00:00 2001 From: Richard Fan Date: Tue, 9 Jun 2026 09:39:21 -0700 Subject: [PATCH 1/2] fix(session-server): restore lock across Phase 1+2+3, preempt DELETE via cancellation channel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #31 split session.lock around the SGLang proxy call to keep DELETE non-blocking while a chat was in flight. That split-lock window let two same-session writers commit interleaved assistant turns; the Phase-3 stale-update guard at sessions.py:241 silently dropped the second commit via `logger.warning(... skipping state update); return backend.build_proxy_response(result)` — the caller still received the 200 SGLang response and appended a phantom assistant record to its trajectory, while the server's session state didn't reflect that turn. The next chat then hit a cursor mismatch. Restore the original "one lock around Phase 1+2+3" invariant and move DELETE preemption from "no lock during proxy" to an explicit cancellation channel: chat stores its proxy task on session.current_proxy_task, DELETE cancels it (after setting session.closing), the chat coroutine catches CancelledError and returns 410 Gone (distinct from 404 "never existed" and 409 "retryable conflict"), then releases the lock so DELETE can acquire it and tear down. SessionStateConflictError (409) is added as defense-in-depth: with Phase 1+2+3 under one lock the Phase-3 stale-state branch is unreachable, but a future split-lock regression now surfaces a clear retryable 409 instead of a silently dropped commit. Two tests in test_session_race_conditions.py encoded the old "DELETE waits for chat to finish" contract and now expect the new "DELETE preempts → 410" behaviour: - test_double_delete_second_returns_404 - test_multiple_chats_queued_then_delete Pre-existing 4-failure baseline (get_or_create_session auto-create-after-delete tests) is unchanged: 5 pass / 4 fail. Split out of #33 per offline review request — #33 keeps the multi-backend orchestration + observability work. Co-Authored-By: Claude Opus 4.7 (1M context) --- miles/rollout/session/linear_trajectory.py | 7 + miles/rollout/session/session_errors.py | 23 ++- miles/rollout/session/sessions.py | 183 +++++++++++------- .../router/test_session_race_conditions.py | 21 +- 4 files changed, 150 insertions(+), 84 deletions(-) diff --git a/miles/rollout/session/linear_trajectory.py b/miles/rollout/session/linear_trajectory.py index fb34b6a0c7..4410cf0ef5 100644 --- a/miles/rollout/session/linear_trajectory.py +++ b/miles/rollout/session/linear_trajectory.py @@ -49,6 +49,13 @@ class LinearTrajectory: lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) closing: bool = field(default=False, repr=False, compare=False) + # DELETE preemption channel: the chat-completions handler stashes its + # in-flight proxy task here so a concurrent DELETE can cancel it without + # waiting on the lock. The chat coroutine catches CancelledError and + # returns 410 Gone. Cleared in the chat handler's finally. + current_proxy_task: "asyncio.Task | None" = field( + default=None, repr=False, compare=False + ) messages: list[dict[str, Any]] = field(default_factory=list) records: list[SessionRecord] = field(default_factory=list) trajectory_token_ids: list[list[int]] = field(default_factory=list) diff --git a/miles/rollout/session/session_errors.py b/miles/rollout/session/session_errors.py index 30e6784384..d028a7bc6c 100644 --- a/miles/rollout/session/session_errors.py +++ b/miles/rollout/session/session_errors.py @@ -3,10 +3,11 @@ Hierarchy --------- SessionError (base) -├── SessionNotFoundError → 404 session does not exist -├── MessageValidationError → 400 messages structure/content invalid -├── TokenizationError → 500 TITO tokenizer / prefix mismatch -└── UpstreamResponseError → 502 SGLang response invalid or unexpected +├── SessionNotFoundError → 404 session does not exist +├── MessageValidationError → 400 messages structure/content invalid +├── SessionStateConflictError → 409 Phase-3 commit guard fired (defense-in-depth) +├── TokenizationError → 500 TITO tokenizer / prefix mismatch +└── UpstreamResponseError → 502 SGLang response invalid or unexpected """ @@ -32,6 +33,20 @@ class MessageValidationError(SessionError): status_code: int = 400 +class SessionStateConflictError(SessionError): + """Raised as a defensive assert when the Phase-3 commit guard fires. + + With the lock-restored chat flow (lock held through Phase 1+2+3), this + branch is unreachable in practice — no other writer can mutate the + session while the proxy is in flight. We keep the guard + 409 surface + as a defense in depth: if a future change reintroduces a split-lock + window, callers see a clear retryable conflict instead of a silently + dropped commit that would corrupt the trajectory's accumulated state. + """ + + status_code: int = 409 + + class TokenizationError(SessionError): """Raised when TITO tokenization invariants are violated. diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index f3de329816..36011fc3e8 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -11,6 +11,7 @@ from miles.rollout.session.session_errors import ( SessionError, SessionNotFoundError, + SessionStateConflictError, TokenizationError, UpstreamResponseError, ) @@ -98,10 +99,17 @@ async def get_session(session_id: str): @app.delete("/sessions/{session_id}") async def delete_session(session_id: str): + # Lock-restored chat flow holds session.lock through Phase 1+2+3. + # Preserve DELETE preemption by cancelling the in-flight proxy task + # (the cancellation channel stored on the session). The cancelled + # chat returns 410 to its caller, then releases the lock so DELETE + # can acquire it and remove the session. session = registry.get_session(session_id) if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") session.closing = True + if session.current_proxy_task is not None: + session.current_proxy_task.cancel() logger.debug( f"[session-server] DELETE waiting for lock: session={session_id} lock_locked={session.lock.locked()}" ) @@ -117,12 +125,23 @@ async def delete_session(session_id: str): async def chat_completions(request: Request, session_id: str): """Proxy a chat completion through SGLang with TITO token tracking. - Flow: prepare pretokenized input_ids (lock held briefly) → inject - SGLang flags → proxy to backend (NO lock) → validate response → - update trajectory checkpoint (lock held briefly) → append session record. - - The lock is NOT held during the slow proxy call to avoid blocking - DELETE/other operations when the agent disconnects mid-request. + Flow: ALL three phases run under ``session.lock`` — + Phase 1: prepare pretokenized input_ids + inject SGLang flags + Phase 2: proxy to backend (lock held; cancellation channel via + ``session.current_proxy_task`` lets DELETE preempt) + Phase 3: validate response, update trajectory checkpoint, record + + Holding the lock through Phase 2 eliminates the cursor-mismatch race + that the split-lock design (lock-Phase-1 → unlock-Phase-2 → relock- + Phase-3) introduced: two same-session writers could both reach + Phase 3, the stale-state guard silently dropped the second commit + while still returning 200, and the caller appended a phantom + assistant record to its trajectory. + + DELETE preemption is preserved via the cancellation channel: + ``delete_session`` cancels ``session.current_proxy_task``; the chat + coroutine catches ``CancelledError`` and returns 410 Gone, then + releases the lock so DELETE can acquire it. """ _inflight_chat["count"] += 1 try: @@ -130,12 +149,12 @@ async def chat_completions(request: Request, session_id: str): if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") - # --- Phase 1: prepare request (lock held briefly) --- async with session.lock: # Double-check: session may have been marked closing while waiting for lock. if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") + # --- Phase 1: prepare request --- body = await request.body() request_body = json.loads(body) if body else {} @@ -171,82 +190,99 @@ async def chat_completions(request: Request, session_id: str): body = json.dumps(request_body).encode() expected_num_assistant = session.num_assistant - # --- lock released here --- - - # --- Phase 2: proxy to SGLang (NO lock held) --- - result = await backend.do_proxy(request, "v1/chat/completions", body=body) - - # If SGLang returned a non-200 error (e.g. 400 for context too long), - # pass it through to the agent without recording — the agent can retry - # or handle the error. - if result["status_code"] != 200: - # Rollback failures indicate corrupted prefix-cache state in SGLang. - # Retry once without pretokenized input_ids so SGLang processes the - # request from scratch instead of attempting prefix continuation. - error_body = result.get("response_body") or b"" - if isinstance(error_body, bytes): - error_body = error_body.decode("utf-8", errors="replace") - if ( - result["status_code"] == 400 - and "rollback failed" in error_body.lower() - and "input_ids" in request_body - ): - logger.warning( - "SGLang rollback failed for session %s, retrying without prefix continuation", - session_id, - ) - request_body.pop("input_ids", None) - retry_body = json.dumps(request_body).encode() - result = await backend.do_proxy(request, "v1/chat/completions", body=retry_body) - if result["status_code"] != 200: - return backend.build_proxy_response(result) - else: - return backend.build_proxy_response(result) - response = json.loads(result["response_body"]) + # --- Phase 2: proxy to SGLang (lock held; cancellation channel) --- + session.current_proxy_task = asyncio.create_task( + backend.do_proxy(request, "v1/chat/completions", body=body) + ) + try: + result = await session.current_proxy_task + except asyncio.CancelledError: + # DELETE preempted this request. Surface 410 Gone to the caller + # so litellm/harbor treats it as a definite session-closed signal + # (distinct from 404 "never existed" and 409 "retryable conflict"). + return JSONResponse(status_code=410, content={"error": "session closing"}) + finally: + session.current_proxy_task = None + + # If SGLang returned a non-200 error (e.g. 400 for context too long), + # pass it through to the agent without recording — the agent can retry + # or handle the error. + if result["status_code"] != 200: + # Rollback failures indicate corrupted prefix-cache state in SGLang. + # Retry once without pretokenized input_ids so SGLang processes the + # request from scratch instead of attempting prefix continuation. + error_body = result.get("response_body") or b"" + if isinstance(error_body, bytes): + error_body = error_body.decode("utf-8", errors="replace") + if ( + result["status_code"] == 400 + and "rollback failed" in error_body.lower() + and "input_ids" in request_body + ): + logger.warning( + "SGLang rollback failed for session %s, retrying without prefix continuation", + session_id, + ) + request_body.pop("input_ids", None) + retry_body = json.dumps(request_body).encode() + session.current_proxy_task = asyncio.create_task( + backend.do_proxy(request, "v1/chat/completions", body=retry_body) + ) + try: + result = await session.current_proxy_task + except asyncio.CancelledError: + return JSONResponse(status_code=410, content={"error": "session closing"}) + finally: + session.current_proxy_task = None + if result["status_code"] != 200: + return backend.build_proxy_response(result) + else: + return backend.build_proxy_response(result) - choice = response.get("choices", [{}])[0] + response = json.loads(result["response_body"]) - meta_info = choice.get("meta_info") - if not isinstance(meta_info, dict) or "output_token_logprobs" not in meta_info: - raise UpstreamResponseError( - "meta_info and output_token_logprobs must be in choice (requires logprobs=True)" - ) - assistant_message = choice.get("message", {}) - if assistant_message.get("content") is None: - raise UpstreamResponseError( - "assistant message content is None, when tool call parser failed SGLang should still return " - "an empty content rather than None. Please check your modified SGLang version." - ) + choice = response.get("choices", [{}])[0] - prompt_token_ids = choice.get("prompt_token_ids") - output_token_logprobs = meta_info["output_token_logprobs"] - completion_tokens = meta_info["completion_tokens"] - - actual_output_logprobs_len = len(output_token_logprobs) - if actual_output_logprobs_len != completion_tokens: - raise UpstreamResponseError( - "invalid chat completion response: " - f"len(output_token_logprobs)={actual_output_logprobs_len} " - f"!= completion_tokens={completion_tokens}. " - f"Please check whether you use the correct SGLang branch which has fix the tokenizer batch decode issue." - ) + meta_info = choice.get("meta_info") + if not isinstance(meta_info, dict) or "output_token_logprobs" not in meta_info: + raise UpstreamResponseError( + "meta_info and output_token_logprobs must be in choice (requires logprobs=True)" + ) + assistant_message = choice.get("message", {}) + if assistant_message.get("content") is None: + raise UpstreamResponseError( + "assistant message content is None, when tool call parser failed SGLang should still return " + "an empty content rather than None. Please check your modified SGLang version." + ) - completion_token_ids = [t[1] for t in output_token_logprobs] + prompt_token_ids = choice.get("prompt_token_ids") + output_token_logprobs = meta_info["output_token_logprobs"] + completion_tokens = meta_info["completion_tokens"] + + actual_output_logprobs_len = len(output_token_logprobs) + if actual_output_logprobs_len != completion_tokens: + raise UpstreamResponseError( + "invalid chat completion response: " + f"len(output_token_logprobs)={actual_output_logprobs_len} " + f"!= completion_tokens={completion_tokens}. " + f"Please check whether you use the correct SGLang branch which has fix the tokenizer batch decode issue." + ) - # --- Phase 3: update state (lock held briefly) --- - async with session.lock: - if session.closing: - logger.warning(f"Session {session_id} closed during proxy, skipping state update") - return backend.build_proxy_response(result) + completion_token_ids = [t[1] for t in output_token_logprobs] + # --- Phase 3: update state (still under the same lock) --- + # Defensive assert: with the lock held through Phase 1+2+3, no + # other writer can mutate num_assistant. If this fires, some + # future change has reintroduced a split-lock window — surface + # 409 instead of silently dropping the commit (the + # cursor-mismatch class of incident under the old design). if session.num_assistant != expected_num_assistant: - logger.warning( - f"Session {session_id} state changed during proxy " + raise SessionStateConflictError( + f"session {session_id} state changed during proxy " f"(expected num_assistant={expected_num_assistant}, " - f"got {session.num_assistant}), skipping state update" + f"got {session.num_assistant})" ) - return backend.build_proxy_response(result) await asyncio.to_thread( session.update_pretokenized_state, @@ -266,7 +302,6 @@ async def chat_completions(request: Request, session_id: str): response=response, ) session.append_record(record) - # --- lock released here --- return backend.build_proxy_response(result) finally: diff --git a/tests/fast/router/test_session_race_conditions.py b/tests/fast/router/test_session_race_conditions.py index 95c6a69cba..0ece86806c 100644 --- a/tests/fast/router/test_session_race_conditions.py +++ b/tests/fast/router/test_session_race_conditions.py @@ -321,7 +321,10 @@ def process_fn(prompt: str) -> ProcessResult: d1 = delete_1.result(timeout=30.0) d2 = delete_2.result(timeout=30.0) - assert chat_resp.status_code == 200 + # Lock-restored chat flow: DELETE preempts the in-flight chat via the + # cancellation channel; the cancelled chat returns 410 Gone. (Under + # the prior split-lock design the chat would complete with 200.) + assert chat_resp.status_code == 410 # One delete succeeds, the other gets 404 codes = sorted([d1.status_code, d2.status_code]) assert codes == [204, 404], f"Expected [204, 404], got {codes}" @@ -385,12 +388,18 @@ def process_fn(prompt: str) -> ProcessResult: assert delete_resp.status_code == 204 - # At least one chat must succeed (the one holding the lock when - # delete arrived). Others may get 200 (acquired lock before - # closing) or 404 (saw closing=True). No 500s allowed. + # Lock-restored chat flow: the in-flight chat at DELETE time gets + # cancelled via the proxy-task channel (→ 410). Queued chats then + # acquire the lock, see closing=True, and return 404. (Under the + # prior split-lock design the in-flight chat would have completed + # with 200 before DELETE acquired the lock.) No 500s allowed. status_codes = [r.status_code for r in results] - assert all(c in (200, 404) for c in status_codes), f"Unexpected status codes: {status_codes}" - assert 200 in status_codes, f"Expected at least one 200, got {status_codes}" + assert all(c in (200, 404, 410) for c in status_codes), ( + f"Unexpected status codes: {status_codes}" + ) + assert 200 in status_codes or 410 in status_codes, ( + f"Expected at least one in-flight chat (200 or 410), got {status_codes}" + ) def test_rapid_create_chat_delete_cycles(self): """Rapidly create, chat, and delete sessions to stress the lifecycle. From a114e9da1148ead6653cd119717f3657ca09f47b Mon Sep 17 00:00:00 2001 From: Richard Fan Date: Tue, 9 Jun 2026 14:40:27 -0700 Subject: [PATCH 2/2] style: pre-commit auto-fixes (ruff + isort + black) Apply formatter fixes flagged by CI's pre-commit run. Three files inside this PR's scope had style-fixes (linear_trajectory.py, test_session_race_conditions.py, plus an unused-import sweep) and seven pre-existing style-debt files elsewhere in the repo that pre-commit's --all-files mode picked up. All changes are formatter output, no semantic edits. Same lint-fix bundle applied to PR #33 in commit e1b93a02. Co-Authored-By: Claude Opus 4.7 (1M context) --- miles/rollout/session/linear_trajectory.py | 4 +--- tests/fast/router/test_session_race_conditions.py | 10 ++++------ 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/miles/rollout/session/linear_trajectory.py b/miles/rollout/session/linear_trajectory.py index 4410cf0ef5..c8eb823473 100644 --- a/miles/rollout/session/linear_trajectory.py +++ b/miles/rollout/session/linear_trajectory.py @@ -53,9 +53,7 @@ class LinearTrajectory: # in-flight proxy task here so a concurrent DELETE can cancel it without # waiting on the lock. The chat coroutine catches CancelledError and # returns 410 Gone. Cleared in the chat handler's finally. - current_proxy_task: "asyncio.Task | None" = field( - default=None, repr=False, compare=False - ) + current_proxy_task: "asyncio.Task | None" = field(default=None, repr=False, compare=False) messages: list[dict[str, Any]] = field(default_factory=list) records: list[SessionRecord] = field(default_factory=list) trajectory_token_ids: list[list[int]] = field(default_factory=list) diff --git a/tests/fast/router/test_session_race_conditions.py b/tests/fast/router/test_session_race_conditions.py index 0ece86806c..b6e2e46b2f 100644 --- a/tests/fast/router/test_session_race_conditions.py +++ b/tests/fast/router/test_session_race_conditions.py @@ -394,12 +394,10 @@ def process_fn(prompt: str) -> ProcessResult: # prior split-lock design the in-flight chat would have completed # with 200 before DELETE acquired the lock.) No 500s allowed. status_codes = [r.status_code for r in results] - assert all(c in (200, 404, 410) for c in status_codes), ( - f"Unexpected status codes: {status_codes}" - ) - assert 200 in status_codes or 410 in status_codes, ( - f"Expected at least one in-flight chat (200 or 410), got {status_codes}" - ) + assert all(c in (200, 404, 410) for c in status_codes), f"Unexpected status codes: {status_codes}" + assert ( + 200 in status_codes or 410 in status_codes + ), f"Expected at least one in-flight chat (200 or 410), got {status_codes}" def test_rapid_create_chat_delete_cycles(self): """Rapidly create, chat, and delete sessions to stress the lifecycle.