refactor(grpc): replace PreparationOutput struct with typed enum

[CatherineSue]

Description

Problem

PreparationOutput is a flat struct with 9 fields, many of which are always None for certain request types. For example, Completion/Generate/Embedding never use processed_messages, tool_constraints, filtered_request, or harmony_* fields — yet they must set them all to None. This wastes memory (the enum is sized to its largest field, ChatCompletionRequest at ~1448 bytes) and provides no type-level guarantee that required fields are present.

Solution

Replace the struct with a 6-variant enum (Chat, Messages, Completion, Generate, Embedding, Harmony), where each variant carries only the fields relevant to its pipeline. Add helper methods (token_ids(), routing_text()) for cross-variant access used by worker selection and simple request builders.

Changes

• context.rs: PreparationOutput struct → enum with 6 variants + helper methods
• 7 preparation stages: construct their specific variant, eliminating boilerplate None fields
• 6 request_building stages: pattern-match destructuring for type-specific fields; helpers for simple cases
• worker_selection: routing_text() replaces if harmony_mode { selection_text } else { original_text }
• completion/response_processing: routing_text() replaces original_text.as_deref()

Key improvements:

• processed_messages is non-optional in Chat/Messages (type-guaranteed, removes runtime .ok_or_else() check)
• Chat/Messages no longer clone original_text (routing_text() borrows from processed_messages.text)
• Harmony request_building uses .take() + move instead of .as_ref() + clone
• Box<ChatCompletionRequest> in Chat/Harmony variants reduces enum size for Completion/Generate/Embedding paths
• debug_assert + graceful gRPC error replaces unreachable!() for variant checks

Test Plan

All existing tests pass (1015 tests across workspace). This is a purely mechanical refactor — no behavioral changes.

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

• Refactor
  • Restructured internal data model for request preparation across multiple request types (chat, messages, completion, generate, embedding, and harmony). Improved data handling and encapsulation through new accessor methods, ensuring more explicit type-safety and cleaner code organization throughout the request processing pipeline.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 1efab3de-34c4-4b30-9577-643e9e5a1924

📥 Commits

Reviewing files that changed from the base of the PR and between 8f92774 and 6a9240d.

📒 Files selected for processing (15)

• model_gateway/src/routers/grpc/common/stages/worker_selection.rs
• model_gateway/src/routers/grpc/context.rs
• model_gateway/src/routers/grpc/harmony/stages/preparation.rs
• model_gateway/src/routers/grpc/harmony/stages/request_building.rs
• model_gateway/src/routers/grpc/regular/stages/chat/preparation.rs
• model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs
• model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs
• model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/embedding/preparation.rs
• model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs
• model_gateway/src/routers/grpc/regular/stages/generate/preparation.rs
• model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs
• model_gateway/src/routers/grpc/regular/stages/messages/preparation.rs
• model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs

---

📝 Walkthrough

Walkthrough

This pull request refactors the PreparationOutput data structure from a struct with numerous optional fields into a tagged enum with request-type-specific variants (Chat, Messages, Completion, Generate, Embedding, Harmony). The change cascades through multiple preparation and request-building stages, replacing struct field access with enum variant destructuring and introducing accessor methods for common fields.

Changes

Cohort / File(s)	Summary
PreparationOutput Refactoring model_gateway/src/routers/grpc/context.rs	Converted PreparationOutput from a struct with optional fields into an enum with six variants, each containing only applicable fields. Added token_ids(&self) and routing_text(&self) accessor methods to encapsulate per-variant logic for extracting common values.
Harmony Pipeline model_gateway/src/routers/grpc/harmony/stages/preparation.rs, model_gateway/src/routers/grpc/harmony/stages/request_building.rs	Updated Harmony preparation to construct PreparationOutput::Harmony variant with non-optional harmony_messages and harmony_stop_ids. Modified request building to consume and pattern-match the preparation output, switching stop-token injection to unconditional logic with direct extend_from_slice calls.
Chat Pipeline model_gateway/src/routers/grpc/regular/stages/chat/preparation.rs, model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs	Updated chat preparation to use PreparationOutput::Chat variant with Box-wrapped filtered_request. Modified request building to destructure the variant and extract fields directly, eliminating fallible accessor paths.
Messages Pipeline model_gateway/src/routers/grpc/regular/stages/messages/preparation.rs, model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs	Changed messages preparation to store data in PreparationOutput::Messages variant. Updated request building to destructure variant fields instead of accessing optional struct members, adding debug assertions for pipeline invariants.
Completion Pipeline model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs, model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs, model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs	Converted completion preparation to PreparationOutput::Completion variant. Updated request building and response processing to use routing_text() accessor method instead of direct original_text field access.
Generate Pipeline model_gateway/src/routers/grpc/regular/stages/generate/preparation.rs, model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs	Modified generate preparation to construct PreparationOutput::Generate variant. Updated request building to use routing_text() and token_ids() accessor methods for deriving request parameters.
Embedding Pipeline model_gateway/src/routers/grpc/regular/stages/embedding/preparation.rs, model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs	Updated embedding preparation to use PreparationOutput::Embedding variant. Modified request building to derive original_text and token_ids from accessor methods instead of direct field cloning.
Worker Selection model_gateway/src/routers/grpc/common/stages/worker_selection.rs	Replaced conditional text selection logic with prep.routing_text() accessor and changed token ID access from direct field inspection to using prep.token_ids() method, simplifying the worker selection input derivation.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

• feat(gateway): wire Messages API pipeline into gRPC routers #753 — Refactor of PreparationOutput and its accessors affects Messages pipeline stages and context usage wired into gRPC routers.
• refactor(grpc): remove redundant harmony_stop_ids from build_generate_request_from_responses #880 — Both PRs modify Harmony request-building stop-token injection logic and how harmony_stop_ids data flows through the pipeline.
• fix(proto): change logprobs token_ids from int32 to uint32 #354 — Both PRs update call sites to use prep.token_ids() accessor and propagate token ID handling through request-building stages.

Suggested labels

grpc, model-gateway

Suggested reviewers

• key4ng
• slin1237

Poem

🐰 Enums now bloom where structs once lay,
Each pipeline finds its proper way,
No more options scattered wide,
Type safety is our trusted guide,
The routing dances, tokens flow so free!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and specifically describes the main change: refactoring PreparationOutput from a struct to a typed enum across the gRPC routing layer.

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

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch refactor/preparation-output-enum

---

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

[gemini-code-assist]

Code Review

This pull request refactors the PreparationOutput struct into an enum with specific variants for each request type, such as Chat, Messages, and Harmony. This change improves type safety by eliminating optional fields that were only relevant to specific pipelines. The review feedback correctly identifies several compilation issues in the request building stages where owned values were moved inside match arms or where type mismatches occurred between Option types and Arc references.

[claude]

Clean mechanical refactor — verified semantic equivalence across all 6 pipelines, take() safety (no later-stage access to preparation), and correct ownership moves in nested match arms. LGTM.