Skip to content

Classify SessionServer endpoints before proxying - #2061

Open
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/session-server-endpoint-classification
Open

Classify SessionServer endpoints before proxying#2061
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/session-server-endpoint-classification

Conversation

@matrix72c

Copy link
Copy Markdown
Contributor

Summary

This PR makes SessionServer classify requests from the HTTP method and normalized endpoint before inspecting or mutating their JSON envelopes.

It fixes Anthropic /v1/messages/count_tokens requests being mistaken for model generations because they contain prompt-shaped fields such as messages, and it closes the catch-all proxy path that could otherwise forward unrelated requests to LMDeploy management endpoints.

This endpoint-policy change was split out of #2059 so that proxy routing and trace-reliability behavior can be reviewed independently.

Endpoint behavior

Request Behavior Training trace
POST /v1/messages Forward as Anthropic generation Yes, when requested
POST /v1/chat/completions Forward as OpenAI generation Yes, when requested
POST /v1/responses Return the existing native 501 response No
POST /v1/messages/count_tokens Transparently forward, removing only a legacy SessionServer-owned session_id No
HEAD /api/hello Return 204 locally No
Any other method or endpoint Return native-format 404 locally No

In particular, GET /v1/models and LMDeploy management routes such as /terminate, /update_weights, /sleep, and /wakeup never reach the worker.

Implementation

  • Add a small private classifier based on method + normalized path.
    • Leading slashes are normalized.
    • Query strings do not participate in classification.
    • Known endpoints may appear behind a base-path prefix.
  • Separate generation processing from auxiliary forwarding before request hooks run.
    • Only generation endpoints run on_request, response cleaning, stream parsing, and on_response.
    • count_tokens preserves messages, system, tools, and other client fields. Only the compatibility-only session_id field is removed.
  • Handle HEAD /api/hello, unsupported methods, and unsupported endpoints locally before opening an upstream connection.
  • Log rejected requests using only the method and normalized path; bodies, headers, query strings, and credentials are not logged.

The classifier is intentionally private and local to SessionServer; this PR does not add a shared enum, dataclass, or cross-project abstraction.

Tests

Added focused classifier and local aiohttp fake-upstream tests covering:

  • leading slash, base-path prefix, query string, and incorrect method handling;
  • Messages, Chat Completions, Responses, token counting, and hello endpoints;
  • transparent token-count request/response forwarding with no trace hooks;
  • unchanged generation hook execution; and
  • local rejection of unknown and LMDeploy management endpoints with zero upstream calls.

Validation performed on this topic alone:

pytest -q tests/rl/test_session_server_endpoints.py
6 passed

ruff check xtuner/v1/rl/rollout/session_server.py tests/rl/test_session_server_endpoints.py
All checks passed

git diff --check upstream/main...fix/session-server-endpoint-classification

The topic also merges cleanly with #2059 in the local integration branch; the combined endpoint and trace suites pass all 12 tests.

Non-goals

  • Trace recovery after malformed successful responses or early client disconnects remains in Preserve SessionServer traces across response failures #2059.
  • No LMDeploy code or management route is changed.
  • Model discovery remains disabled.
  • Claude Code WebFetch, plugin-marketplace, and general tool network policy are outside this proxy path and are unchanged.

@matrix72c
matrix72c force-pushed the fix/session-server-endpoint-classification branch 3 times, most recently from 623f513 to d80dd84 Compare September 1, 2026 13:00
@matrix72c
matrix72c force-pushed the fix/session-server-endpoint-classification branch from d80dd84 to 6165c7b Compare September 1, 2026 13:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant