[None][feat] Add more disagg conversation ID headers support

[reasonsolo]

Summary by CodeRabbit

• New Features

  • Improved conversation identification mechanism with multi-header support and priority-based selection for session headers. Replaces previous single-header limitation, providing enhanced flexibility and robustness across different deployment configurations and request patterns.

• Tests

  • Added comprehensive unit and integration test coverage for conversation ID extraction functionality.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change extends conversation ID extraction in the OpenAI disaggregated server to support multiple session headers with a defined priority order, replacing exclusive reliance on a single header. A new test module validates the extraction logic across various header configurations.

Changes

Cohort / File(s)	Summary
Server Implementation tensorrt_llm/serve/openai_disagg_server.py	Added _CONVERSATION_ID_HEADERS class constant and updated _extract_conversation_id helper to iterate through multiple headers in priority order, extracting the first present header and populating disaggregated_params.conversation_id.
Test Module tests/unittest/disaggregated/test_openai_disagg_server.py	New unittest module with comprehensive tests for _extract_conversation_id method, validating multi-header support, case-insensitive matching, pre-existing body field preservation, and conditional population based on disaggregated_params presence.
Test List Updates tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/test-db/l0_a10.yml	Added new test module references to QA and integration test selection lists.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is entirely composed of the template with all sections left blank—no explanation of the issue, solution, test coverage, or rationale was provided by the author.	Fill in the Description and Test Coverage sections with details about why conversation ID header support was expanded and which tests validate the changes.

✅ Passed checks (3 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly describes the main change: adding support for multiple conversation ID headers in the disaggregated server, which matches the core implementation change.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/disaggregated/test_openai_disagg_server.py (1)

26-57: ⚡ Quick win

Add edge-case coverage for empty and missing headers.

Please add cases for:

1. higher-priority header present but empty while lower-priority header is non-empty, and
2. no supported headers present (assert request remains unchanged).
   This locks down fallback and no-op behavior.

As per coding guidelines: “Coverage expectations: Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix.”

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/disaggregated/test_openai_disagg_server.py` around lines 26 -
57, Extend test_extract_conversation_id_from_headers to cover two edge cases:
add a case where a higher-priority header (e.g., "X-Session-ID") is present but
an empty string while a lower-priority header (e.g., "X-Correlation-ID" or
"x-session-affinity") is non-empty and assert that
OpenAIDisaggServer._extract_conversation_id selects the non-empty lower-priority
value on the CompletionRequest created in the test (use _raw_request(headers)
and assert request.disaggregated_params.conversation_id equals the
lower-priority value); and add a case where none of the supported headers are
present and assert that the CompletionRequest remains unchanged
(request.disaggregated_params is None) after calling _extract_conversation_id.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/serve/openai_disagg_server.py`:
- Around line 190-196: The loop that picks the conversation_id currently treats
empty string values as valid; update the logic in
OpenAIDisaggServer._CONVERSATION_ID_HEADERS iteration (where header_conv_id is
set via raw_req.headers.get) to ignore None or empty/whitespace-only header
values (e.g., check header_conv_id is not None and header_conv_id.strip() != "")
before breaking, so it will continue to lower-priority headers and only return
when a non-empty conversation_id is found.

---

Nitpick comments:
In `@tests/unittest/disaggregated/test_openai_disagg_server.py`:
- Around line 26-57: Extend test_extract_conversation_id_from_headers to cover
two edge cases: add a case where a higher-priority header (e.g., "X-Session-ID")
is present but an empty string while a lower-priority header (e.g.,
"X-Correlation-ID" or "x-session-affinity") is non-empty and assert that
OpenAIDisaggServer._extract_conversation_id selects the non-empty lower-priority
value on the CompletionRequest created in the test (use _raw_request(headers)
and assert request.disaggregated_params.conversation_id equals the
lower-priority value); and add a case where none of the supported headers are
present and assert that the CompletionRequest remains unchanged
(request.disaggregated_params is None) after calling _extract_conversation_id.

🪄 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: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7c37fe27-5384-4f53-aa81-b80db9874eef

📥 Commits

Reviewing files that changed from the base of the PR and between f4563d5 and f8dd09d.

📒 Files selected for processing (4)

• tensorrt_llm/serve/openai_disagg_server.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_a10.yml
• tests/unittest/disaggregated/test_openai_disagg_server.py

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #46346 [ run ] triggered by Bot. Commit: e1b390f Link to invocation

[tensorrt-cicd]

PR_Github #46346 [ run ] completed with state ABORTED. Commit: e1b390f

Link to invocation