[Feat]: finished by finish reason

[scarlet25151]

Pull Request Description

This PR fixes the missing finish_reason handling in PD disaggregation for TensorRT-LLM.

Previously, the gateway always proceeded to the decode phase after a synchronous TRT-LLM prefill request, even when the prefill response had already completed generation. This could break short outputs such as single-token answers (yes/no, numbers, etc.), because the workflow incorrectly assumed a decode hop was still required.

This change adds a TRT-LLM-specific short-circuit for synchronous prefill responses:

• If finish_reason is a completed state (that is, neither "length" nor "not_finished"), the gateway returns the prefill result immediately and skips decode.
• If the request is streaming, the gateway does not short-circuit and preserves the existing streaming behavior.
• The short-circuit is scoped only to TensorRTLLM so other engines, including vLLM, keep their original behavior.

In addition, this PR includes:

• A fix for an immediate-response handling compile issue in gateway.go
• A regression test covering TRT-LLM prefill completion short-circuit behavior
• A regression test ensuring vLLM does not short-circuit on finish_reason
• Defensive nil-context handling in request cleanup paths to avoid panics for invalid requests (for example, unknown model names)
• A mock TRT-LLM PD prefill adjustment in the development app so PD e2e tests continue to exercise the decode path when intended

Related Issues

Resolves: #2013

Important: Before submitting, please complete the description above and review the checklist below.

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

• [Bug]: Corrections to existing functionality
• [CI]: Changes to build process or CI pipeline
• [Docs]: Updates or additions to documentation
• [API]: Modifications to aibrix's API or interface
• [CLI]: Changes or additions to the Command Line Interface
• [Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

• PR title includes appropriate prefix(es)
• Changes are clearly explained in the PR description
• New and existing tests pass successfully
• Code adheres to project style and best practices
• Documentation updated to reflect changes (if applicable)
• Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

[gemini-code-assist]

Code Review

This pull request introduces a short-circuiting mechanism for prefill requests within the PD disaggregation algorithm. It enables the gateway to return an immediate response to the client if a specific finish reason is detected during the prefill phase, effectively bypassing subsequent processing steps. Key changes include modifications to the routing context, HTTP execution flow, and metrics handling, as well as the addition of utility functions for constructing Envoy-compatible immediate responses and new unit tests. I have no feedback to provide.

[scarlet25151]

@varungup90 please help to review

[varungup90]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a short-circuiting mechanism for TRT-LLM prefill requests that complete early (e.g., with a 'stop' finish reason), allowing the gateway to return an immediate response instead of proceeding to a decode phase. It also includes significant hardening of the RoutingContext and associated routing algorithms by adding nil-pointer checks to prevent panics during request cleanup or when handling immediate responses. Feedback focuses on ensuring that the stream context is properly reset when the pooled RoutingContext is deleted to avoid race conditions, and addressing a gap where short-circuited requests currently bypass token accounting and quota metrics.

[gemini-code-assist]

When st.routerCtx is deleted and set to nil, st.ctx should also be reset to the original stream context (e.g., the context from srv.Context()). Since RoutingContext objects are pooled and reused, continuing to use a deleted context in st.ctx can lead to race conditions or unexpected behavior in subsequent iterations of the processing loop (for instance, in preRecvCheck or if srv.Recv() is called again). Additionally, a similar cleanup (setting st.routerCtx = nil and resetting st.ctx) should be performed after the call to HandleResponseBody when st.completed is true, to avoid stale pointers in the processState.

[gemini-code-assist]

Short-circuited prefill responses return an immediate response to the client without proceeding to the response processing phase or calling AddRequestCount. Consequently, tokens generated during the prefill phase (which are present in the routingCtx.ImmediateResponse.Body) are not accounted for in metrics, logs, or user quotas. This could allow users to bypass token-based rate limits for requests that complete during prefill. Consider extracting the usage information from the response body and reporting it (e.g., via s.requestEndHelper) before returning the immediate response.

[varungup90]

📝 Review Comments for PR #2061

🔴 Critical & Correctness

• Fix for routingCtx.Stream (Silent Bug): In gateway_req_body.go:532, routingCtx.Stream was previously never set. This meant shouldShortCircuitPrefillResponse would always default to stream=false, potentially breaking streaming requests. This is a significant bug fix that should be explicitly called out in the Pull Request Description as a separate bullet point.
• Insecure Blocklist Approach: The shouldShortCircuitFinishReason currently uses a blocklist (excluding "", length, not_finished). This is brittle; if TRT-LLM introduces new intermediate states (e.g., tool_calls), the gateway will incorrectly short-circuit. Suggestion: Switch to an allowlist of known completed states like stop, content_filter, and eos_token.
• Header Priority Inversion: In util.go:689-706, extraHeaders (intended for gateway-injected data like X-Request-ID) are being overwritten by response headers from TRT-LLM. Please document if this precedence is intentional, as it prevents the gateway from guaranteed ID injection.

🏗️ Design & Architecture

• Nil-Receiver Proliferation: Adding if r == nil to 10+ methods in router_context.go is highly defensive and may mask upstream bugs where the context is missing.
  • Recommendation: Instead of making every method nil-safe, implement a single helper: GetRoutingContext(ctx) *RoutingContext.
• Refactor requestEndHelper: The nil-check logic in gateway_rsp_body.go:597-612 is fragmented. It mixes nil-safe receiver calls with explicit nil checks for direct field access. Please group direct field accesses inside a single nil-guard block to improve readability.

🧪 Testing & Minor Improvements

• Test Coverage: TestTensorRTPrefillFinishReasonShortCircuit tests the prefill logic directly but misses the full Route() → HandleRequestBody pipeline. The path where the target pod is set to nil (pd_disaggregation.go:342) needs explicit integration coverage.
• Test Consistency: TestVLLMPrefillFinishReasonDoesNotShortCircuit should be converted to a table-driven test to match the TRT-LLM test style.
• CI Reliability: TestPowerOfTwoRouter_DoneRequestCountNilContext is currently gated by testing.Short(). Since it only checks for panics and doesn't strictly require Redis, the gate should be removed to ensure it runs in all CI environments.

[Jeffwan]

@scarlet25151 could you double check varun's command and high priority items from gemini?

[scarlet25151]

@scarlet25151 could you double check varun's command and high priority items from gemini?

@Jeffwan I discuss with varun, and there is some of a easier way to implement, currently I'm testing on new implementation to avoid regression on vllm, will push it later