[Fix] Classify SessionServer endpoints and fail closed on trace errors - #2061
Merged
YifanHe-ailab merged 6 commits intoSep 18, 2026
Merged
YifanHe-ailab merged 6 commits into
YifanHe-ailab merged 6 commits into
Conversation
matrix72c
force-pushed
the
fix/session-server-endpoint-classification
branch
5 times, most recently
from
September 2, 2026 03:36
dd8954e to
cd15002
Compare
matrix72c
changed the title
matrix72c
changed the title
matrix72c
changed the title
matrix72c
force-pushed
the
fix/session-server-endpoint-classification
branch
from
September 4, 2026 09:12
5a8433a to
76eaaba
Compare
matrix72c
force-pushed
the
fix/session-server-endpoint-classification
branch
from
September 7, 2026 03:52
76eaaba to
7ea06f4
Compare
YifanHe-ailab
self-requested a review
September 14, 2026 07:31
Collaborator
|
@claude review |
Collaborator
|
@claude review |
Contributor
|
Claude: Summary本 PR 为 Main Flowchart after this PRflowchart TD
A[_handle_request] --> B{_is_generation_endpoint<br/>method + normalized path}
B -->|否| C[仅剥离 session_id<br/>不跑 on_request / clean_data]
B -->|是| D[filter tools / on_request<br/>计算 trace_enabled]
C --> E[转发 upstream]
D --> E
E --> F{2xx?}
F -->|否| G[原样回传,不清洗不 trace]
F -->|是| H{is_stream?}
H -->|否| I[clean_data + 解析 JSON object]
H -->|是| J[逐行 clean_data<br/>skip_blank_after_done 状态机<br/>扣留 DONE]
J --> K[_parse_stream_response]
K -->|_UpstreamResponseError| L[按上游错误放行<br/>依据 raw 中是否含 DONE 决定补发]
K -->|其他异常 / None| M[session_error_msg]
I --> N[on_response]
K --> N
N --> O{client_alive?}
O -->|是| P[写 error 或 DONE,再 write_eof]
O -->|否| Q[不写 error、不 write_eof<br/>直接返回 response]
style B fill:#cde4ff
style C fill:#cde4ff
style F fill:#cde4ff
style J fill:#ffe3b3
style L fill:#ffb3b3
style Q fill:#ffb3b3
核心原理实现与单测
公开 Interface 的线性业务流程评估
单测建议
其他 Issues
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Comment on lines
939
to
943
| for choice in event.get("choices", []): | ||
| if choice.get("finish_reason") == "error": | ||
| raise RuntimeError( | ||
| raise _UpstreamResponseError( | ||
| f"Upstream SSE choice finished with error: {json.dumps(event, ensure_ascii=False)}" | ||
| ) |
Contributor
There was a problem hiding this comment.
Claude: [测试] finish_reason == "error" 改抛 _UpstreamResponseError 后无用例覆盖。该分支在上游已发 [DONE] 时使 skip_done=False 并补发 [DONE],客户端看到语法完整但未入 trace store 的流,与 fail-closed 目标相悖。
ProduceBatchResult impact: 该样本不会被标为 error,本应计入 leftover_failed 的 turn 被当作 COMPLETED 计入 produced_samples / raw_rewards_count。
建议补一个 upstream SSE 以 finish_reason:"error" + data: [DONE] 结尾的真实用例,断言不补发 [DONE]。
Comment on lines
789
to
805
| if is_stream: | ||
| try: | ||
| if session_error_msg: | ||
| error_payload = _error_payload(fmt, session_error_msg) | ||
| error_line = "data: " + json.dumps(error_payload, ensure_ascii=False) + "\n\n" | ||
| # Anthropic SSE carries a named ``event: error`` line; the recorder keys on the ``data:`` JSON | ||
| # either way, but a real Anthropic SDK consumer needs the event name. | ||
| if fmt == FMT_ANTHROPIC: | ||
| error_line = "event: error\n" + error_line | ||
| await response.write(error_line.encode("utf-8")) | ||
| skip_done = True | ||
| if not skip_done: | ||
| await response.write(b"data: [DONE]\n\n") | ||
| await response.write_eof() | ||
| if client_alive: | ||
| if session_error_msg: | ||
| error_payload = _error_payload(fmt, session_error_msg) | ||
| error_line = "data: " + json.dumps(error_payload, ensure_ascii=False) + "\n\n" | ||
| # Anthropic SSE carries a named ``event: error`` line; the recorder keys on the ``data:`` JSON | ||
| # either way, but a real Anthropic SDK consumer needs the event name. | ||
| if fmt == FMT_ANTHROPIC: | ||
| error_line = "event: error\n" + error_line | ||
| await response.write(error_line.encode("utf-8")) | ||
| skip_done = True | ||
| elif trace_response and not skip_done: | ||
| await response.write(b"data: [DONE]\n\n") | ||
| await response.write_eof() | ||
| except (ConnectionError, ClientConnectionResetError): | ||
| pass |
Contributor
There was a problem hiding this comment.
Claude: [正确性] write_eof() 被收进新增的 if client_alive:,下游中途断连时不再调用;改前它始终在 try/except ConnectionError 中被尝试。另外 prepare 失败分支会把一个未 prepare 的 StreamResponse 返回给 aiohttp。建议 write_eof() 移出 client_alive 分支(仍由外层 except 兜底)。
Comment on lines
+203
to
+289
|
|
||
| async def test_disconnect_before_prepare_still_completes_trace(self): | ||
| self.upstream_content_type = "text/event-stream" | ||
| self.upstream_body = self._openai_stream() | ||
|
|
||
| class Request: | ||
| is_proxy = True | ||
| method = "POST" | ||
| match_info = {"path": "v1/chat/completions"} | ||
| query_string = "" | ||
| headers = {"Content-Type": "application/json"} | ||
|
|
||
| async def read(self): | ||
| return json.dumps(TestSessionServerTraceHandling._request_payload(stream=True)).encode() | ||
|
|
||
| class FakeStreamResponse: | ||
| def __init__(self, status, headers): | ||
| self.status = status | ||
| self.headers = headers | ||
| self.prepare_calls = [] | ||
|
|
||
| async def prepare(self, request): | ||
| self.prepare_calls.append(request) | ||
| raise ClientConnectionResetError("downstream closed") | ||
|
|
||
| async def write_eof(self): | ||
| pass | ||
|
|
||
| fake_web = SimpleNamespace( | ||
| StreamResponse=FakeStreamResponse, | ||
| Response=web.Response, | ||
| json_response=web.json_response, | ||
| ) | ||
| with patch.object(session_server_module, "web", fake_web): | ||
| response = await self.session_server._handle_request(Request()) | ||
|
|
||
| self.assertEqual(response.status, 200) | ||
| self.assertEqual(len(response.prepare_calls), 1) | ||
| self.session_server.on_response.assert_awaited_once() | ||
|
|
||
| async def test_disconnect_midstream_still_completes_trace(self): | ||
| self.upstream_content_type = "text/event-stream" | ||
| self.upstream_body = self._openai_stream() | ||
|
|
||
| class Request: | ||
| is_proxy = True | ||
| method = "POST" | ||
| match_info = {"path": "v1/chat/completions"} | ||
| query_string = "" | ||
| headers = {"Content-Type": "application/json"} | ||
|
|
||
| async def read(self): | ||
| return json.dumps(TestSessionServerTraceHandling._request_payload(stream=True)).encode() | ||
|
|
||
| class FakeStreamResponse: | ||
| def __init__(self, status, headers): | ||
| self.status = status | ||
| self.headers = headers | ||
| self.prepare_calls = [] | ||
| self.write_calls = [] | ||
|
|
||
| async def prepare(self, request): | ||
| self.prepare_calls.append(request) | ||
|
|
||
| async def write(self, data): | ||
| self.write_calls.append(data) | ||
| raise ClientConnectionResetError("downstream closed") | ||
|
|
||
| async def write_eof(self): | ||
| pass | ||
|
|
||
| fake_web = SimpleNamespace( | ||
| StreamResponse=FakeStreamResponse, | ||
| Response=web.Response, | ||
| json_response=web.json_response, | ||
| ) | ||
| with patch.object(session_server_module, "web", fake_web): | ||
| response = await self.session_server._handle_request(Request()) | ||
|
|
||
| self.assertEqual(response.status, 200) | ||
| self.assertEqual(len(response.prepare_calls), 1) | ||
| self.assertEqual(len(response.write_calls), 1) | ||
| self.session_server.on_response.assert_awaited_once() | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() |
Contributor
There was a problem hiding this comment.
Claude: [测试] 这两个 disconnect 用例用 SimpleNamespace 整体替换模块级 web,并自造 Request / FakeStreamResponse,assertEqual(response.status, 200) 只是读回假对象自己的属性。真实 aiohttp 下「prepare 失败后返回未 prepare 的 StreamResponse」「write 失败后跳过 write_eof」这两个新行为都没被验证。建议改用真实 TestClient 提前关闭连接来触发。
Comment on lines
+683
to
+727
| try: | ||
| await response.prepare(request) | ||
| except (ConnectionError, ClientConnectionResetError): | ||
| client_alive = False | ||
| else: | ||
| client_alive = True | ||
| skip_blank_after_done = False | ||
| async for line in resp.content: | ||
| # Only retain chunks when we'll actually need to parse | ||
| # them for tracing; evaluate-mode requests skip this | ||
| # so memory does not grow with stream length. | ||
| if trace_enabled: | ||
| if trace_response: | ||
| response_chunks.append(line) | ||
|
|
||
| if request_data is not None and line.startswith(b"data: ") and line.strip() != b"data: [DONE]": | ||
| if ( | ||
| is_success_response | ||
| and clean_data is not None | ||
| and request_data is not None | ||
| and line.startswith(b"data: ") | ||
| and line.strip() != b"data: [DONE]" | ||
| ): | ||
| try: | ||
| text = line.decode("utf-8") | ||
| data = json.loads(text[6:]) | ||
| if clean_data(data): | ||
| if isinstance(data, dict) and not _is_error_payload(data) and clean_data(data): | ||
| line = ("data: " + json.dumps(data) + "\n").encode("utf-8") | ||
| except Exception: | ||
| pass | ||
|
|
||
| # Delay [DONE] only while a training trace still needs to be exported. | ||
| if client_alive and (not trace_enabled or line.strip() != b"data: [DONE]"): | ||
| if trace_response and skip_blank_after_done: | ||
| if not line.strip(): | ||
| continue | ||
| skip_blank_after_done = False | ||
| if line.strip() == b"data: [DONE]": | ||
| if trace_response: | ||
| skip_blank_after_done = True | ||
| continue | ||
| if client_alive: | ||
| try: | ||
| await response.write(line) | ||
| except (ConnectionError, ClientConnectionResetError): | ||
| client_alive = False | ||
|
|
||
| raw_response = b"".join(response_chunks) if trace_enabled else b"" | ||
| raw_response = b"".join(response_chunks) if trace_response else b"" |
Contributor
There was a problem hiding this comment.
Claude: [设计] Files: xtuner/v1/rl/rollout/session_server.py::SessionServer._handle_request
Problem: _handle_request 本次再增约 70 行后达到约 275 行,把 endpoint 策略、body 改写、header 构造、上游转发、SSE 逐行清洗、skip_blank_after_done 的 [DONE] 抑制状态机、trace 解析与错误下发全部内联在同一抽象层级。trace_enabled / trace_active / trace_response / is_success_response / skip_done / skip_blank_after_done / client_alive 七个布尔在同一作用域交叉约束,读者必须同时持有全部状态才能判断某一行是否会被写出;这也正是上面两条 Warning(write_eof 被 client_alive 吞掉、SSE 流内只能靠假 web 模块测试)难以被发现和覆盖的直接原因。
Solution: 把主流程写成同一抽象层级的线性叙述,机制下沉到私有方法:
async def _handle_request(self, request):
policy = self._classify_request(request) # endpoint / fmt / trace 意图
forward = await self._prepare_forward(request, policy) # body 改写 + headers
outcome = await self._forward(forward, policy) # 转发 + 清洗 + 收集 raw_response
trace_error = await self._export_trace(outcome, policy) # 解析 + on_response
return await self._finalize(outcome, trace_error) # DONE / error / write_eof注意不要机械地每步抽一个函数:按 Deletion test,只保留能真正隐藏规则的私有方法——SSE 行改写与 [DONE] 抑制状态机是最值得单独成 Module 的一块(它自身就是一个可独立测试的纯逐行变换),_finalize 则把 client_alive 与收尾写入的不变量集中到一处,write_eof 是否调用的规则就只剩一个位置可改。
Benefits: Leverage 上调用者与测试跨越同一个 Seam,SSE 抑制逻辑可直接用字节序列断言而不必伪造 web 模块;Locality 上 [DONE] 与 write_eof 的规则各自收敛到一处,后续再加 endpoint 或状态码策略时不需要在 275 行主流程里重新推演七个布尔的组合。
YifanHe-ailab
approved these changes
Sep 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SessionServeris a catch-all HTTP proxy, but only two requests represent model generations:POST /v1/messagesandPOST /v1/chat/completions.This change makes endpoint policy explicit before request/response processing and makes trace-enabled generations fail closed when the response cannot be recorded. Auxiliary API calls remain transparent, while the client-visible result stays consistent with the training trace stored by
SessionServer.The response-reliability changes previously proposed in #2059 are consolidated here so endpoint and trace policy are reviewed in one PR.
Problem
The proxy cannot safely infer the endpoint from the JSON body. For example,
POST /v1/messages/count_tokenscontainsmessagesbut is not a model generation. Treating it as one can invokeon_request, inject generation-only fields, filter tools, or invokeon_responseunexpectedly.For trace-enabled generations, malformed JSON, non-object response bodies, incomplete SSE streams, downstream disconnects, or a failing
on_responsehook can leave the trace store without a complete assistant turn. Returning a normal success response in that situation allows the caller to observe a turn that cannot be used for training.Implementation
_is_generation_endpoint(method, path)predicate. It uses the normalized method and path, ignores query strings, and recognizes onlyPOST /v1/messagesandPOST /v1/chat/completionsas generation endpoints./v1/messages/count_tokens,/v1/messages/batches, and future/v1/messages/...paths still receive the correct Anthropic headers._handle_request; noContextVar, wrapper handler, shared enum, or cross-project abstraction is introduced.session_id. Non-object bodies and objects withoutsession_idremain byte-for-byte unchanged whenever possible.2xxresponses never enter the cleaner or trace hooks, and explicit upstream error envelopes are passed through without being wrapped again.[DONE]until parsing andon_responsecomplete. If parsing or the response hook fails, return a native500error or SSE error event without[DONE].Successful generation behavior, the public
SessionServerAPI, and the existing/v1/responses501rejection remain unchanged. No LMDeploy code or management endpoint implementation is changed, andfix/chunk-loss-detached-head-memoryis untouched.Tests
Added focused
aiohttpfake-upstream coverage for:count_tokensand auxiliary endpoint forwarding with no generation hooks;on_response;2xxresponses remaining unchanged;2xxresponses failing closed;[DONE]; andValidation performed: