Skip to content

feat(grpc): add sglang reasoning token usage - #1747

Merged
slin1237 merged 2 commits into
smg-project:mainfrom
Moersity:lixiang/grpc_reasoning_token
Jun 21, 2026
Merged

slin1237 merged 2 commits into
smg-project:mainfrom
Moersity:lixiang/grpc_reasoning_token

Conversation

@Moersity

@Moersity Moersity commented Jun 16, 2026 •

Copy link
Copy Markdown
Contributor

Description

Problem

SGLang gRPC responses did not expose reasoning_tokens usage through SMG, even when SGLang itself was able to return reasoning token counts.

The gRPC path also did not forward the require_reasoning flag to the SGLang engine, so SGLang would not count reasoning tokens for thinking-enabled requests.

Solution

This PR propagates reasoning token usage through the SGLang gRPC stack:

  • Adds reasoning_tokens to SGLang gRPC stream and completion responses.
  • Adds require_reasoning to SGLang gRPC generate requests.
  • Forwards require_reasoning from the gateway only when thinking/reasoning is effectively enabled.
  • Exposes reasoning_tokens in regular generate/chat responses and streaming usage.
  • Keeps non-SGLang backends unaffected by returning 0 reasoning tokens for unsupported engines.

Changes

  • Added reasoning_tokens fields to SGLang gRPC proto response messages.
  • Added require_reasoning to SGLang gRPC generate request messages.
  • Forwarded require_reasoning in the Python SGLang gRPC servicer.
  • Extracted reasoning_tokens from SGLang batch outputs and returned them through gRPC responses.
  • Exposed reasoning token usage in SMG router usage accounting.
  • Added reasoning token handling for streaming and non-streaming paths.
  • Updated Go binding proto parsing for the new response fields.
  • Introduced request build option structs to keep builder APIs clippy-clean after adding reasoning support.

Test Plan

Validated against an actual SGLang gRPC-backed SMG deployment with thinking-capable models.

## start sglang server with grpc mode
SGLANG_ENABLE_SPEC_V2=1 sglang serve --host 0.0.0.0 \
    --model-path /tmp/GLM-5.1-FP8 \
    --tp 8 --trust-remote-code \
    --reasoning-parser glm45 --tool-call-parser glm47 \
    --speculative-algorithm EAGLE --speculative-num-steps 3 \
    --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \
    --mem-fraction-static 0.85 --enable-cache-report \
    --grpc-mode 
## start smg 
cargo run --bin smg  launch --worker-urls grpc://127.0.0.1:30000 \
    --port 8000 --reasoning-parser glm45   \
    --tool-call-parser glm47_moe

Tested cases:

Validated against an actual SGLang gRPC-backed SMG deployment.

Tested cases:

  • Non-streaming chat completion with thinking enabled:
    • Response includes reasoning_content.
    • usage.completion_tokens_details.reasoning_tokens is present and greater than 0.
curl -s http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -H 'Authorization: Bearer EMPTY' \
  -d @- <<EOF
  {
  "model": "/tmp/GLM-5.1-FP8",
  "messages": [
    {
      "role": "system",
      "content": "你是一个严谨的技术助手。下面是一段固定长前缀,用于测试prefix cache命中。请保持回答简洁。固定前缀开始:KV Cache用于缓存Transformer每一层历史token的key/value,避免多轮生成时重复计算历史上下文。Prefix Cache进一步复用多个请求之间相同前缀的KV,从而降低prefill开销,改善TTFT。在LLM推理服务中,如果多个请求共享相同system prompt、工具定义、RAG模板或长上下文前缀,就可以命中前缀缓存。固定前缀结束。"
    },
    {
      "role": "user",
      "content": "为什么prefix cache可以降低TTFT?"
    }
  ],
  "temperature": 0
}
EOF

image
  • Streaming chat completion with thinking enabled:
    • Streaming response completes normally.
    • Final usage chunk includes non-zero reasoning_tokens.
data: {"id":"chatcmpl-019ecf88-94a3-76f2-a7a8-e4551c9fb35d","object":"chat.completion.chunk","created":1781598295,"model":"/tmp/GLM-5.1-FP8","system_fingerprint":"default","choices":[],"usage":{"prompt_tokens":131,"completion_tokens":453,"total_tokens":584,"prompt_tokens_details":{"cached_tokens":128},"completion_tokens_details":{"reasoning_tokens":397}}}
  • Chat completion with thinking explicitly disabled:
    • require_reasoning is not forwarded as true.
    • reasoning_tokens is absent or 0.
curl -s http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -H 'Authorization: Bearer EMPTY' \
  -d @- <<EOF
  {
  "model": "/tmp/GLM-5.1-FP8",
  "messages": [
    {
      "role": "system",
      "content": "你是一个严谨的技术助手。下面是一段固定长前缀,用于测试prefix cache命中。请保持回答简洁。固定前缀开始:KV Cache用于缓存Transformer每一层历史token的key/value,避免多轮生成时重复计算历史上下文。Prefix Cache进一步复用多个请求之间相同前缀的KV,从而降低prefill开销,改善TTFT。在LLM推理服务中,如果多个请求共享相同system prompt、工具定义、RAG模板或长上下文前缀,就可以命中前缀缓存。固定前缀结束。"
    },
    {
      "role": "user",
      "content": "为什么prefix cache可以降低TTFT?"
    }
  ],
  "stream": true,
  "stream_options": {
    "include_usage": true
  },
  "chat_template_kwargs": {"enable_thinking": false},
  "temperature": 0
}
EOF

## result 
data: {"id":"chatcmpl-019ecf8a-48f8-7581-8635-6764c850ea93","object":"chat.completion.chunk","created":1781598406,"model":"/tmp/GLM-5.1-FP8","system_fingerprint":"default","choices":[],"usage":{"prompt_tokens":131,"completion_tokens":101,"total_tokens":232,"prompt_tokens_details":{"cached_tokens":128}}}

data: [DONE]
  • Non-SGLang backend compatibility:
    • Existing gRPC backends continue to work.
    • Unsupported backends report reasoning_tokens=0.
Checklist
  • cargo +nightly fmt passes
  • cargo clippy --all-targets --all-features -- -D warnings passes
  • (Optional) Documentation updated
  • (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

Summary by CodeRabbit

  • New Features
    • Added reasoning token counting and reporting in both streaming and completed generation responses, including usage aggregation.
    • Introduced a require_reasoning flag derived from chat content and tokenizer settings, and propagated it through gRPC request building for chat and messages.
    • Added unified request options for multimodal inputs, tool-call constraints, and reasoning behavior.
  • Bug Fixes
    • Parsing now safely defaults missing/invalid reasoning token values to 0.
  • Tests
    • Added coverage for request/response reasoning propagation.
  • Chores
    • Version bumps: Python gRPC client 0.4.10 → 0.4.11 and Python servicer 0.5.5 → 0.5.6.

@github-actions github-actions Bot added dependencies Dependency updates grpc gRPC client and router changes protocols Protocols crate changes model-gateway Model gateway crate changes labels Jun 16, 2026
@coderabbitai

coderabbitai Bot commented Jun 16, 2026 •

Copy link
Copy Markdown

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds require_reasoning as a boolean field to the SGLang gRPC GenerateRequest proto and reasoning_tokens counters to GenerateStreamChunk and GenerateComplete. These fields are wired through the Python servicer, Rust gRPC client (via a new SglangGenerateRequestOptions struct), model gateway (via GenerateRequestBuildOptions), Go FFI bindings, and all streaming/non-streaming response-formatting paths. A new chat_requires_reasoning utility determines whether reasoning should be enabled based on tokenizer thinking-key values and toggle policies.

Changes

Reasoning Token Support End-to-End

Layer / File(s) Summary
Proto contract: require_reasoning and reasoning_tokens fields
crates/grpc_client/proto/sglang_scheduler.proto
Adds bool require_reasoning = 18 to GenerateRequest, uint32 reasoning_tokens = 9 to GenerateStreamChunk, and uint32 reasoning_tokens = 12 to GenerateComplete.
Tokenizer utilities: chat_requires_reasoning logic
bindings/golang/src/utils.rs
Adds Tokenizer trait imports and introduces chat_requires_reasoning function that extracts thinking-related values from chat_template_kwargs using the tokenizer's thinking key, then applies tokenizer thinking_toggle() policies to determine whether reasoning token counting should be enabled.
Rust SGLang scheduler library: SglangGenerateRequestOptions and builder refactor
crates/grpc_client/src/sglang_scheduler.rs, crates/grpc_client/src/lib.rs, crates/grpc_client/python/pyproject.toml
Defines SglangGenerateRequestOptions grouping multimodal_inputs, tool_call_constraint, and require_reasoning; updates build_generate_request_from_chat and build_generate_request_from_messages method signatures; derives require_reasoning from JSON body or defaults to false in plain/responses/completion builder paths; re-exports struct from crate root; bumps Python package to 0.4.11; adds tokio tests verifying require_reasoning forwarding.
Proto response wrappers and metadata: reasoning_tokens accessors and field
model_gateway/src/routers/grpc/proto_wrapper.rs, crates/protocols/src/generate.rs
Adds reasoning_tokens() methods to both ProtoGenerateStreamChunk and ProtoGenerateComplete returning the SGLang field value or 0 for other backends; adds optional reasoning_tokens field to GenerateMetaInfo.
Go FFI utilities and bindings: chat_requires_reasoning FFI export
bindings/golang/src/preprocessor.rs, bindings/golang/internal/ffi/preprocessor.go, bindings/golang/src/proto_parse.rs
Exports C FFI function sgl_chat_requires_reasoning_with_tokenizer accepting request JSON and tokenizer handle; implements Go wrapper ChatRequiresReasoningWithTokenizer marshaling data across C boundary; Go proto parsing reads reasoning_tokens from JSON responses in parse_chunk and parse_complete, converting to u32 and defaulting to 0.
Model Gateway client: GenerateRequestBuildOptions and GrpcClient method refactor
model_gateway/src/routers/grpc/client.rs
Introduces GenerateRequestBuildOptions struct with multimodal_inputs, tool_constraints, and require_reasoning; updates build_chat_request and build_messages_request method signatures; routes require_reasoning into SglangGenerateRequestOptions for SGLang backend and sources multimodal/tool-constraint args from options for vLLM, TRT-LLM, MLX, and TokenSpeed.
Request building stages: chat, messages, Harmony, and Go bindings
model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs, model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs, model_gateway/src/routers/grpc/harmony/stages/request_building.rs, bindings/golang/src/client.rs, bindings/golang/src/policy.rs
Chat stage derives require_reasoning from tokenizer and chat_template_kwargs; messages stage derives it from messages_request.thinking; Harmony stage forces it to false; Go FFI client and policy bindings compute require_reasoning via ffi.ChatRequiresReasoningWithTokenizer and construct SglangGenerateRequestOptions with explicit field mapping.
Python gRPC servicer: require_reasoning input and reasoning_tokens output
grpc_servicer/smg_grpc_servicer/sglang/servicer.py, grpc_servicer/smg_grpc_servicer/sglang/request_manager.py, grpc_servicer/pyproject.toml
HealthCheck sets require_reasoning=False; _convert_generate_request passes require_reasoning from gRPC request into TokenizedGenerateReqInput; request_manager extracts per-request reasoning_tokens into locals during batch-output assembly; streaming chunk and completion builders populate reasoning_tokens from meta_info; package bumped to 0.5.6 with dependency on smg-grpc-proto >=0.4.11.
Go client integration: require_reasoning determination and proto parsing
bindings/golang/internal/grpc/client_grpc.go, bindings/golang/src/grpc_converter.rs
CreateChatCompletionStream calls ffi.ChatRequiresReasoningWithTokenizer to determine require_reasoning flag and wires it into GenerateRequest; protoToJSON includes reasoning_tokens fields from Chunk and Complete proto responses; grpc_converter chains with_reasoning_tokens into final usage builder.
Response formatting and streaming: reasoning_tokens aggregation and propagation
model_gateway/src/routers/grpc/common/response_formatting.rs, model_gateway/src/routers/grpc/regular/streaming.rs, model_gateway/src/routers/grpc/regular/processor.rs
build_usage sums reasoning_tokens across completions via with_reasoning_tokens; streaming chat adds per-index map populated on Complete events and flushed into final SSE Usage; generate streaming meta_info includes reasoning_tokens in chunk and complete events for single-stream and dual/logprob paths; non-streaming processor populates reasoning_tokens in GenerateMetaInfo.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

  • lightseekorg/smg#1031: Introduces the underlying utils::extract_thinking_from_kwargs and utils::should_mark_reasoning_started logic for computing require_reasoning; this PR now uses those utilities directly in the request-building stages.
  • lightseekorg/smg#744: Adds the initial messages request-building stage and build_generate_request_from_messages backend plumbing; this PR refactors both to use SglangGenerateRequestOptions and propagate reasoning fields through that path.
  • lightseekorg/smg#337: Introduces the shared gRPC response-formatting helpers (Usage::with_reasoning_tokens(...) and common usage construction); this PR extends them to aggregate reasoning tokens across all response paths.

Suggested reviewers

  • slin1237
  • key4ng
  • CatherineSue

Poem

🐇 Hop hop, a new field appears,
require_reasoning calms our fears!
reasoning_tokens count each thought,
Through proto, Rust, and Python wrought.
From chunk to complete, the numbers flow —
A bunny's math puts on a show! 🧮

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately describes the main objective: adding reasoning token usage support to the SGLang gRPC stack, which is the central focus of all file changes across proto definitions, gRPC servicer, gateway, and Go bindings.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for SGLang reasoning tokens across the gRPC client, servicer, and model gateway. It introduces a require_reasoning option for generation requests and propagates reasoning_tokens through streaming chunks and completion responses. Feedback on the changes suggests using getattr when accessing reasoning_tokens on BatchTokenIDOutput in Python to maintain backward compatibility with older SGLang versions and avoid potential runtime crashes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread grpc_servicer/smg_grpc_servicer/sglang/request_manager.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f7216b6fdc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread grpc_servicer/smg_grpc_servicer/sglang/servicer.py
Comment thread bindings/golang/src/proto_parse.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
model_gateway/src/routers/grpc/regular/streaming.rs (1)

2362-2365: ⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoff

Missing reasoning_tokens tracking in completion streaming endpoint.

The completion streaming endpoint (process_completion_streaming_chunks) does not track or expose reasoning_tokens, while the chat streaming (lines 213, 494, 567, 573-574) and generate streaming (lines 791-792, 821, 981-982, 1025) endpoints both include reasoning token accounting. This creates an inconsistency where users calling /v1/completions with thinking-enabled models will not receive reasoning_tokens in the usage statistics.

Add reasoning token tracking to the completion streaming endpoint to maintain feature parity:

  1. At line 2365: Add let mut total_reasoning = 0u32;
  2. At line 2484: Add total_reasoning = total_reasoning.max(complete.reasoning_tokens());
  3. At line 2622: Chain .with_reasoning_tokens(total_reasoning) after .with_cached_tokens(total_cached)

Also applies to: 2480-2484, 2613-2623

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/routers/grpc/regular/streaming.rs` around lines 2362 -
2365, The process_completion_streaming_chunks function is missing
reasoning_tokens tracking that exists in other streaming endpoints (chat
streaming and generate streaming). Add reasoning token tracking by: (1)
initializing a new mutable variable `total_reasoning` as 0u32 alongside the
existing `total_prompt` and `total_cached` variable declarations, (2)
accumulating reasoning tokens from each completion chunk using
`total_reasoning.max(complete.reasoning_tokens())` in the same section where
`total_cached` is being accumulated, and (3) chaining
`.with_reasoning_tokens(total_reasoning)` onto the response builder after the
existing `.with_cached_tokens(total_cached)` call to expose reasoning tokens in
the usage statistics.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@model_gateway/src/routers/grpc/regular/streaming.rs`:
- Around line 2362-2365: The process_completion_streaming_chunks function is
missing reasoning_tokens tracking that exists in other streaming endpoints (chat
streaming and generate streaming). Add reasoning token tracking by: (1)
initializing a new mutable variable `total_reasoning` as 0u32 alongside the
existing `total_prompt` and `total_cached` variable declarations, (2)
accumulating reasoning tokens from each completion chunk using
`total_reasoning.max(complete.reasoning_tokens())` in the same section where
`total_cached` is being accumulated, and (3) chaining
`.with_reasoning_tokens(total_reasoning)` onto the response builder after the
existing `.with_cached_tokens(total_cached)` call to expose reasoning tokens in
the usage statistics.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: cb663fb4-28d5-4017-980a-5e3a21fc6718

📥 Commits

Reviewing files that changed from the base of the PR and between ca22340 and c7efa4b.

📒 Files selected for processing (19)
  • bindings/golang/src/client.rs
  • bindings/golang/src/policy.rs
  • bindings/golang/src/proto_parse.rs
  • crates/grpc_client/proto/sglang_scheduler.proto
  • crates/grpc_client/python/pyproject.toml
  • crates/grpc_client/src/lib.rs
  • crates/grpc_client/src/sglang_scheduler.rs
  • crates/protocols/src/generate.rs
  • grpc_servicer/pyproject.toml
  • grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
  • grpc_servicer/smg_grpc_servicer/sglang/servicer.py
  • model_gateway/src/routers/grpc/client.rs
  • model_gateway/src/routers/grpc/common/response_formatting.rs
  • model_gateway/src/routers/grpc/harmony/stages/request_building.rs
  • model_gateway/src/routers/grpc/proto_wrapper.rs
  • model_gateway/src/routers/grpc/regular/processor.rs
  • model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs
  • model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs
  • model_gateway/src/routers/grpc/regular/streaming.rs

@mergify

mergify Bot commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

Hi @Moersity, the DCO sign-off check has failed. All commits must include a Signed-off-by line.

To fix existing commits:

# Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-lease

To sign off future commits automatically:

  • Use git commit -s every time, or
  • VSCode: enable Git: Always Sign Off in Settings
  • PyCharm: enable Sign-off commit in the Commit tool window

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ea87d900a0

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread crates/grpc_client/src/sglang_scheduler.rs Outdated
@Moersity
Moersity force-pushed the lixiang/grpc_reasoning_token branch from ea87d90 to 7caf6fd Compare June 16, 2026 09:09

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@bindings/golang/src/client.rs`:
- Around line 220-224: The require_reasoning field is hardcoded to false in two
locations, preventing proper reasoning-token accounting for thinking-enabled
requests. In bindings/golang/src/client.rs at lines 220-224, compute the
require_reasoning value from the incoming chat request and tokenizer using the
same decision logic applied in gateway request stages, then pass the computed
value to SglangGenerateRequestOptions instead of hardcoding false. Apply the
identical computed require_reasoning behavior in bindings/golang/src/policy.rs
at lines 634-638 for the multi-worker request path, replacing the hardcoded
false value with the same computed logic.

In `@crates/grpc_client/src/sglang_scheduler.rs`:
- Around line 226-230: The require_reasoning extraction and propagation logic in
the sglang_scheduler.rs module lacks direct test coverage. Add focused unit
tests that verify: (1) parsing of the require_reasoning boolean from request
JSON body.other, including edge cases where the value is true, false, or a
non-boolean type (should default to false), and (2) the forwarding of
require_reasoning from chat/messages options into the generated proto message.
These tests should exercise the and_then().unwrap_or(false) pattern shown in the
require_reasoning extraction to ensure correct behavior across all input
scenarios.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: f964cea5-4bbc-45dc-b55e-e067d9580f51

📥 Commits

Reviewing files that changed from the base of the PR and between c7efa4b and 7caf6fd.

📒 Files selected for processing (20)
  • bindings/golang/src/client.rs
  • bindings/golang/src/grpc_converter.rs
  • bindings/golang/src/policy.rs
  • bindings/golang/src/proto_parse.rs
  • crates/grpc_client/proto/sglang_scheduler.proto
  • crates/grpc_client/python/pyproject.toml
  • crates/grpc_client/src/lib.rs
  • crates/grpc_client/src/sglang_scheduler.rs
  • crates/protocols/src/generate.rs
  • grpc_servicer/pyproject.toml
  • grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
  • grpc_servicer/smg_grpc_servicer/sglang/servicer.py
  • model_gateway/src/routers/grpc/client.rs
  • model_gateway/src/routers/grpc/common/response_formatting.rs
  • model_gateway/src/routers/grpc/harmony/stages/request_building.rs
  • model_gateway/src/routers/grpc/proto_wrapper.rs
  • model_gateway/src/routers/grpc/regular/processor.rs
  • model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs
  • model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs
  • model_gateway/src/routers/grpc/regular/streaming.rs

Comment thread bindings/golang/src/client.rs
Comment thread crates/grpc_client/src/sglang_scheduler.rs

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7caf6fdfa6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread bindings/golang/src/client.rs Outdated
@Moersity
Moersity force-pushed the lixiang/grpc_reasoning_token branch from 7caf6fd to 412f79f Compare June 16, 2026 09:29

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 412f79f411

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".


// Ask SGLang scheduler to count reasoning tokens.
// Keep false unless the request explicitly enables reasoning/thinking.
bool require_reasoning = 18;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Regenerate Go SDK protos for reasoning fields

The schema adds require_reasoning/reasoning_tokens, but the checked-in pure Go SDK path still uses the stale generated proto under bindings/golang/internal/proto (I checked with rg and there are no RequireReasoning or ReasoningTokens accessors), and bindings/golang/internal/grpc/client_grpc.go builds GenerateRequest/serializes GenerateResponse directly through those types. The fresh path I checked is separate from the fixed Rust FFI call sites, so Go SDK streaming calls still cannot request reasoning and will drop the returned counter before postprocessing; please regenerate/update the Go proto and conversion code along with this schema change.

Useful? React with 👍 / 👎.

@Moersity
Moersity force-pushed the lixiang/grpc_reasoning_token branch from 412f79f to 404b653 Compare June 16, 2026 09:59

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 404b6535ac

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines 589 to +591
Usage::from_counts(prompt_tokens, completion_tokens)
.with_cached_tokens(complete.cached_tokens),
.with_cached_tokens(complete.cached_tokens)
.with_reasoning_tokens(complete.reasoning_tokens),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Expose reasoning usage in Go SDK structs

When SGLang returns non-zero complete.reasoning_tokens, this now serializes completion_tokens_details.reasoning_tokens into the JSON chunk, but the Go SDK typed path still unmarshals usage into bindings/golang/client.go's Usage struct, which only has prompt/completion/total fields; CreateChatCompletion and the multi-client path then copy that struct and silently drop the new detail. This affects Go callers using the typed APIs rather than raw RecvJSON; please add matching completion-token detail fields and preserve them while accumulating usage.

Useful? React with 👍 / 👎.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
model_gateway/src/routers/grpc/regular/streaming.rs (1)

2612-2627: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Missing reasoning_tokens in completions streaming usage chunk.

The completions streaming endpoint tracks total_cached and includes it in the final usage chunk (line 2622), but does not track or include reasoning_tokens. Chat streaming (lines 567-574) correctly tracks reasoning tokens per index and includes them via .with_reasoning_tokens(total_reasoning). Since both endpoints use the same Usage structure from openai_protocol::common::Usage, completions should also track and report reasoning tokens.

📊 Proposed fix to add reasoning_tokens tracking

Add reasoning token tracking at the top of the function, similar to total_cached:

     let mut total_prompt = 0u32;
     let mut total_cached = 0u32;
+    let mut total_reasoning = 0u32;
     let mut total_completion = CompletionTokenTracker::new();

Track reasoning tokens when processing Complete messages:

                 ProtoResponseVariant::Complete(complete) => {
                     let index = complete.index();
                     total_prompt = total_prompt.max(complete.prompt_tokens());
                     total_cached = total_cached.max(complete.cached_tokens());
+                    total_reasoning = total_reasoning.max(complete.reasoning_tokens());
                     total_completion.record_complete(&complete);

Include reasoning tokens in the final usage chunk:

         if include_usage {
             let usage_chunk = CompletionStreamResponse {
                 id: request_id.clone(),
                 object: "text_completion".to_string(),
                 created,
                 choices: vec![],
                 model: model.clone(),
                 system_fingerprint: system_fingerprint.map(String::from),
                 usage: Some(
                     Usage::from_counts(total_prompt, total_completion.total())
                         .with_cached_tokens(total_cached)
+                        .with_reasoning_tokens(total_reasoning),
                 ),
             };
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/routers/grpc/regular/streaming.rs` around lines 2612 -
2627, The completions streaming endpoint is not tracking reasoning_tokens in the
final usage chunk, while the chat streaming endpoint correctly implements this.
Add reasoning token tracking to match the pattern used for total_cached:
initialize a variable to accumulate reasoning tokens at the top of the function
containing the CompletionStreamResponse creation, update this variable when
processing Complete messages (similar to how total_cached is updated), and then
chain a call to .with_reasoning_tokens(total_reasoning) after the
.with_cached_tokens(total_cached) call on the Usage::from_counts() chain to
include reasoning tokens in the usage chunk sent via
Self::format_completion_sse_into.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/grpc_client/proto/sglang_scheduler.proto`:
- Line 249: Add an inline comment to the reasoning_tokens field in the proto
file to document its purpose. The comment should clarify that this field
represents the count of reasoning or thinking tokens used in the completed
response, similar to how other token count fields are documented in the same
message. Place the comment directly above or inline with the uint32
reasoning_tokens = 12 field definition to improve code maintainability and
readability.
- Line 223: Add an inline comment to the reasoning_tokens field in the proto
file to document that it represents cumulative reasoning/thinking token usage.
Place the comment directly above or next to the field definition to clarify its
semantics and maintain consistency with the documentation style of adjacent
token fields like prompt_tokens, completion_tokens, and cached_tokens.

In `@model_gateway/src/routers/grpc/client.rs`:
- Around line 45-50: The GenerateRequestBuildOptions struct lacks documentation
explaining its purpose and the purpose of its fields. Add a doc comment above
the struct definition explaining that it consolidates request-building
parameters, then add individual doc comments for each field (multimodal_inputs,
tool_constraints, and require_reasoning) describing what each field represents
and how it's used in request building.

---

Outside diff comments:
In `@model_gateway/src/routers/grpc/regular/streaming.rs`:
- Around line 2612-2627: The completions streaming endpoint is not tracking
reasoning_tokens in the final usage chunk, while the chat streaming endpoint
correctly implements this. Add reasoning token tracking to match the pattern
used for total_cached: initialize a variable to accumulate reasoning tokens at
the top of the function containing the CompletionStreamResponse creation, update
this variable when processing Complete messages (similar to how total_cached is
updated), and then chain a call to .with_reasoning_tokens(total_reasoning) after
the .with_cached_tokens(total_cached) call on the Usage::from_counts() chain to
include reasoning tokens in the usage chunk sent via
Self::format_completion_sse_into.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 26d62cbf-1dfc-428c-afdf-a5f7d16aae2d

📥 Commits

Reviewing files that changed from the base of the PR and between 7caf6fd and 404b653.

⛔ Files ignored due to path filters (3)
  • bindings/golang/internal/proto/common.pb.go is excluded by !**/*.pb.go
  • bindings/golang/internal/proto/sglang_scheduler.pb.go is excluded by !**/*.pb.go
  • bindings/golang/internal/proto/sglang_scheduler_grpc.pb.go is excluded by !**/*.pb.go
📒 Files selected for processing (24)
  • bindings/golang/internal/ffi/preprocessor.go
  • bindings/golang/internal/grpc/client_grpc.go
  • bindings/golang/src/client.rs
  • bindings/golang/src/grpc_converter.rs
  • bindings/golang/src/policy.rs
  • bindings/golang/src/preprocessor.rs
  • bindings/golang/src/proto_parse.rs
  • bindings/golang/src/utils.rs
  • crates/grpc_client/proto/sglang_scheduler.proto
  • crates/grpc_client/python/pyproject.toml
  • crates/grpc_client/src/lib.rs
  • crates/grpc_client/src/sglang_scheduler.rs
  • crates/protocols/src/generate.rs
  • grpc_servicer/pyproject.toml
  • grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
  • grpc_servicer/smg_grpc_servicer/sglang/servicer.py
  • model_gateway/src/routers/grpc/client.rs
  • model_gateway/src/routers/grpc/common/response_formatting.rs
  • model_gateway/src/routers/grpc/harmony/stages/request_building.rs
  • model_gateway/src/routers/grpc/proto_wrapper.rs
  • model_gateway/src/routers/grpc/regular/processor.rs
  • model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs
  • model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs
  • model_gateway/src/routers/grpc/regular/streaming.rs

Comment thread crates/grpc_client/proto/sglang_scheduler.proto
Comment thread crates/grpc_client/proto/sglang_scheduler.proto
Comment thread model_gateway/src/routers/grpc/client.rs

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f2017a3ea5

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines +100 to +104
let require_reasoning = ctx.tokenizer_arc().is_some_and(|tokenizer| {
utils::should_mark_reasoning_started(
utils::extract_thinking_from_kwargs(
chat_request.chat_template_kwargs.as_ref(),
tokenizer.as_ref(),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Set require_reasoning for always-reasoning models

For chat requests whose tokenizer reports ThinkingToggle::None because the model always reasons (the enum documents DeepSeek R1 as that case), this predicate returns false, so the SGLang request is sent with require_reasoning=false. Those models can still emit reasoning that the parser separates, but SGLang's scheduler only fills reasoning_tokens when the request bit is true, so usage stays zero for DeepSeek-R1-style chat requests; please cover the analogous Messages/Go call sites that use the same predicate too.

Useful? React with 👍 / 👎.

@Moersity

Copy link
Copy Markdown
Contributor Author

Sorry for the earlier CI noise. I missed running the full pre-commit checks locally before pushing.

I’ve fixed the issue, rerun the checks locally, and pushed the updated commit. This should be ready for review again. Thanks for taking another look.

@Moersity

Copy link
Copy Markdown
Contributor Author

@coderabbitai review


// Ask SGLang scheduler to count reasoning tokens.
// Keep false unless the request explicitly enables reasoning/thinking.
bool require_reasoning = 18;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a requirement for sglang to start count for reasoning tokens? Why do they need such a field instead of turning it on by default? Is it due to performance concerns?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, if not set, reasoning tokens will always zero, see sglang, This field is describe: Note: For native API, as a work-around, you need to set require_reasoningargument toTrue to ensure the model will think before generating the structured output. It's not required for chat-completion API.

@mergify

mergify Bot commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Hi @Moersity, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch:

git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease

@mergify mergify Bot added the needs-rebase PR has merge conflicts that need to be resolved label Jun 18, 2026
@mergify mergify Bot removed the needs-rebase PR has merge conflicts that need to be resolved label Jun 19, 2026
@Moersity
Moersity force-pushed the lixiang/grpc_reasoning_token branch from 847fe8c to e1f430d Compare June 19, 2026 14:50

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e1f430dbb1

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

top_logprobs_num: body.logprobs.unwrap_or(0) as i32,
return_hidden_states: body.return_hidden_states,
stream: body.stream,
require_reasoning: false,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Honor require_reasoning on completion requests

For /v1/completions callers that pass the engine-specific require_reasoning: true extra field (the request type preserves unknown fields in CompletionRequest.other, and this path already forwards SGLang-specific constraints such as json_schema), the new proto bit is always sent as false. SGLang only fills reasoning_tokens when this request flag is true, so structured/reasoning completion calls will still report zero reasoning usage even though the generate/chat paths now propagate the flag; please read the boolean from body.other here as the plain generate builder does.

Useful? React with 👍 / 👎.

Signed-off-by: Moersity <lixiang0417.cq@gmail.com>
@Moersity
Moersity force-pushed the lixiang/grpc_reasoning_token branch from e1f430d to 02bca87 Compare June 19, 2026 15:02
@Moersity

Copy link
Copy Markdown
Contributor Author

Hi everyone, this PR is ready for review. Please let me know if there are any additional requirements or improvements I should address to get this approved and merged. Looking forward to your feedback!

@slin1237
slin1237 merged commit 7cb64bf into smg-project:main Jun 21, 2026
46 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dependencies Dependency updates grpc gRPC client and router changes model-gateway Model gateway crate changes protocols Protocols crate changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants