[sgl-router] refactor - layout BucketResolver, Bucket, EngineGroup and implement PowerOfTwo (#40241)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-20 16:57:07 -07:00
parent 4a9dc5c4af
commit aedda8377e
15 changed files with 2566 additions and 4 deletions
@@ -0,0 +1,575 @@
 # Engine selection: target design
 This document describes the target architecture for engine selection. It defines
 responsibilities and behavior; the interface and configuration examples are
 sketches, not a specification of the current API or CLI. Implementation status is
 listed at the end.
 ## Principles
 1. **Order compatible buckets by token length first.** `BucketResolver` returns
   all matching buckets, smallest capacity first, without inspecting workers or policies.
 2. **The bucket owns plain versus PD engine selection.** `Bucket::pick_engines`
   calls its one plain group or both prefill and decode groups, returning a
   complete selection. Both PD engines come from that same bucket.
 3. **An engine group owns membership and its policy.** `EngineGroup::pick`
   filters live workers by model, health, stage, and membership, invokes the
   policy, and validates that its result belongs to the exact candidate set.
 4. **A policy owns selection, fallback, admission, and its state dependencies.**
   Construction injects the load, KV, or affinity handles it needs. Request-time
   arguments contain facts and candidates. A policy cannot choose another bucket
   or cross a PD role boundary.
 5. **The handler owns bucket fallback and dispatch.** It calls `pick_engines`
   on each bucket and dispatches only after a complete selection succeeds.
   Missing candidates or admission rejection advance to the next bucket, where
   all required engines are selected again. Invalid signals or policy results stop routing.
 An engine is represented by `Worker`. Registry role labels remain authoritative
 when filtering candidates. Groups reference worker IDs rather than owning live
 workers. Group policy instances are reused across requests.
 ### Data ownership
 | Type | Owns |
 | --- | --- |
 | `BucketResolver` | A model's bucket collection, length filtering, and ordering |
 | `Bucket` | ID, length constraints, rank, groups, and complete plain/PD engine selection |
 | `EngineGroup` | Engine membership, attached policy, and engine selection |
 | `WorkerRegistry` | Live workers, model membership, health, and role |
 | Request handler | Request preparation, ordered bucket attempts, HTTP errors, and dispatch |
 `worker_ids: None` means every healthy engine serving the requested model and
 role. An explicit empty set means no engines. `EngineGroup::new(policy)` creates
 a catch-all membership group. `BucketGroups::Pd` requires both groups, making
 partial or mixed plain/PD bucket configurations unrepresentable.
 ## 1. Code organization
 ```text
 src/
  buckets.rs                    Bucket resolution and engine-group selection
  policies/
    mod.rs                      Policy contract and construction
    admission.rs                Per-engine acceptance checks
    cache_aware.rs               Cache selection and local/remote adapter
    session_aware.rs             Session selection and fallback
    sticky.rs                   Routing-key selection and fallback
    least_load.rs               Least-load selection
    power_of_two.rs              Pair sampling and stage-aware comparison
    random.rs                   Random selection
    round_robin.rs              Rotation
  state/
    mod.rs                      Shared exports
    kv_events/                  Local index, subscriptions, hashing, wire format
    load_monitor/               Engine reports and router-local request accounting
    affinity_store.rs           Assignments, expiry, and atomic updates
  server/
    app_context.rs              Shared service lifecycle and policy wiring
    routes/chat.rs              Request preparation, PD coordination, dispatch
 ```
 During migration, `buckets_reorg.rs` and `policies_reorg/` implement this design
 beside the live `policies/` path. The layout above is the target after switchover.
 Shared state already lives directly under `src/state/`.
 Dependencies flow from engine groups to policies, and from policies and admission
 to shared state. State does not depend on bucket ordering or concrete policy
 strategies. Small shared helpers are sufficient; no generic score-composition,
 tier executor, or separate selection framework is required.
 ## 2. Responsibilities and request flow
 `BucketResolver::resolve(input_tokens, expected_peak_tokens)` returns an ordered
 list of compatible bucket references (possibly empty), or an invalid-signal error.
 It does not receive a stage or a load view, resolve live engines, or invoke policies.
 The handler iterates this list until a bucket supplies the complete engine selection.
 `Bucket::pick_engines(workers, request)` accepts a `BucketRequest`
 of prepared routing facts, invokes the required groups, and returns `BucketPick`
 (one plain pick or a complete P/D pair). Failures retain their stage. This API
 has no HTTP headers, `AppContext`, or forwarding dependency.
 `EngineGroup::pick(workers, request)` resolves healthy workers for the request's
 model and stage, intersects them with its membership, sorts them by stable ID,
 and invokes its attached policy. It rejects foreign results, including a newly
 allocated worker with the same ID as a candidate.
 ```text
 chat_completions (reorg configured): prepare tokens and expected peak
  |
  v
 BucketResolver::resolve: ordered length-compatible buckets
  |
  v
 For each bucket: bucket.pick_engines(...)
  |
  +-- BucketGroups::Plain
  |     plain.pick() -> one plain engine
  |
  +-- BucketGroups::Pd
        prefill.pick() -> P engine
        decode.pick()  -> D engine from the same bucket
  |
  +-- empty group / admission rejection -> try next bucket (repeat all picks)
  +-- invalid signal / configuration / foreign pick -> return error
  |
  v
 Complete selection -> forward_chat_request: acquire guards, attach PD bootstrap, forward response
 ```
 The handler extracts token facts and header keys once into `BucketRequest`.
 The bucket creates a stage-specific `PickRequest` for each group call, supplying
 its own ID and the role associated with that group. Input length, expected peak,
 token IDs, and session/routing keys pass through. Policies obtain observations
 from their own shared-state handles; buckets and handlers do not provide load,
 KV, or affinity services on each call. PD uses separate group policies, but never
 independently resolves a decode bucket. Decode selection
 failure discards that tentative prefill choice and advances to the next bucket
 on missing candidates or admission rejection. No forwarding guards are acquired
 and no prefill request is sent until both picks in one bucket succeed.
 PD compatibility constraints beyond model and role remain follow-up work.
 There is one endpoint: `POST /v1/chat/completions`. `AppContext::chat_routing`
 chooses its implementation:
 - `ChatRouting::Legacy` (default) uses the existing policies and bucket selector.
 - `ChatRouting::Reorg(HashMap<ModelId, BucketResolver>)` uses the new bucket and
  policy interfaces, with explicit model-specific resolvers.
 Callers set this field before building the router. A missing model in the reorg
 map returns 404, without falling back to legacy routing. This PR adds the
 programmatic configuration switch; CLI/configuration factory construction and
 the remaining production policies remain follow-ups. Power-of-two is implemented
 for explicit attachments; the default serving path remains legacy.
 Both implementations reuse request preparation (including sampling validation
 and tokenization), forwarding, streaming, middleware, and the 32 MiB body limit.
 The reorg implementation requests tokenization for length matching, retaining
 the existing body-size estimate when tokenization is unavailable.
 ## 3. Bucket resolution
 1. Validate that a known expected peak is at least the input length.
 2. Keep buckets whose inclusive input-token range contains the input length.
 3. Check the bucket context capacity against input plus requested output when
   known, or against input length when the output budget is unknown.
 4. Sort by ascending input capacity (the lesser of the input upper bound and
   context capacity). Unbounded capacities sort last. Break ties by ascending
   bucket rank, then ID, and return the entire ordered list.
 5. The handler calls each bucket's `pick_engines` until one supplies its complete
   selection. A failed PD attempt never contributes an engine to a later pair.
 The handler checks addition overflow when computing the expected peak.
 An empty bucket list becomes a 400 `NoMatchingBucket` response. After exhausting
 the list, accumulated admission rejection details produce a selection failure
 (503); if there were no admission rejections, the last unavailable stage produces
 a stage-specific 503. Invalid policy signals/configuration or out-of-candidate
 picks stop the pass immediately with an internal error. Successful engine
 selection ends the pass; forwarding errors do not restart bucket iteration.
 Token ranges and rank belong to the bucket, not its engine groups. For a PD
 bucket, both groups share this one request-length decision. Policy fallback on
 a cache/affinity miss stays within that group's candidates. There is no second
 pass with relaxed admission and no post-policy substitution.
 SLO ordering, cache lookup, and session/sticky policies are follow-ups. Their
 integration must preserve bucket-first selection and the same-bucket PD rule.
 Cross-bucket affinity probing is not part of this interface. Session/routing
 keys still pass through `PickRequest` for policies operating inside the selected
 group; unsupported legacy modes need explicit migration decisions before the
 standard serving path switches.
 ## 4. Policy and admission contracts
 ### Policy
 All policies implement one asynchronous, object-safe `Policy::pick` interface.
 A boxed future supports the remote prefix indexer; local policies can return an
 immediately ready result.
 ```rust
 pub trait Policy: Send + Sync + std::fmt::Debug {
    fn pick<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>>;
 }
 ```
 Policies own their required state handles and read observations locally. Nested
 fallback uses `pick_fallback(engines, request)`, which calls the fallback's `pick`;
 the fallback reads its own state. There is no shared observation context or cache.
 Buckets and HTTP handlers supply only candidates and request facts.
 `Pick` identifies one engine and a selection reason for metrics and tracing.
 `PickRequest` carries model, stage, selected bucket ID, input and optional
 expected peak counts, optional token IDs, and session/routing keys. It contains
 no HTTP body, bucket resolver, state handles, snapshots, or backend configuration.
 | Outcome | Meaning |
 | --- | --- |
 | `Pick` | The chosen engine belongs to the supplied set and passed admission |
 | `NoMatchingBucket` | No bucket supports the requested length |
 | `NoCandidates` | No eligible member or policy selection miss |
 | `NoAdmissibleEngine` | A policy exhausted its candidates; includes rejection reasons |
 | `AdmissionRejected` | The chosen engine failed after-selection admission |
 | `InvalidSignal` / `InvalidConfiguration` | Invalid policy input or configuration |
 | `OutsideCandidates` | Policy returned an engine outside its exact candidate set |
 ### Admission
 Admission evaluates acceptance. It does not rank engines, choose replacements,
 change buckets, or mutate affinity.
 | Check | Acceptance rule |
 | --- | --- |
 | `AllowAll` | Add no acceptance constraint |
 | `CapacityAdmission` | Projected running requests and KV tokens fit reported capacity |
 | `PendingPrefillAdmission` | Waiting uncached tokens plus incoming uncached work fit the budget |
 | `InFlightLimitAdmission` | Router-local in-flight requests are below the limit |
 | `QueueLimitAdmission` | Engine-reported waiting requests are below the limit |
 | `AllOfAdmission` | Every attached check allows the request |
 `EngineAdmission::check(engine, request, load)` checks one engine and returns
 `Allow`, `Reject(reason)`, or an error for invalid inputs. Policies attach the
 checker directly as `Arc<dyn EngineAdmission>`. There is no placement setting,
 filtering wrapper, or before/after API; each policy decides where checking
 belongs in its selection algorithm. The `load` argument is an
 `Option<&EngineReportedWorkerLoad>` retained by the policy for this engine, including
 request counts, token usage, capacity, and the report timestamp. `None` means
 no usable observation, never zero load; each check defines its missing-data
 behavior. Other required state handles belong to the checker.
 Power-of-two first selects an engine, then calls admission exactly once on that
 engine. A rejection returns `AdmissionRejected` to the bucket loop; it does not
 resample, choose the other sampled engine, or run a policy fallback. No candidates
 returns `NoCandidates` without invoking admission. A single candidate is selected
 directly; otherwise two distinct candidates are sampled uniformly, and the one
 with lower stage pressure wins. A complete tie keeps the first sampled engine.
 Power-of-two reuses the existing pure pressure-comparison functions. Plain and
 prefill stages compare estimated prefill queue time when both reports provide it,
 then waiting uncached tokens, waiting requests, and running requests. Decode
 compares waiting requests, running requests, KV usage fraction, then used KV tokens.
 Reported-pressure ties use router-local active requests. If either sampled engine
 lacks a fresh, complete native report with valid capacity, both are compared by
 router-local active requests instead. Basic reports from older publishers are
 still passed to admission when fresh, but do not supply native pressure metrics.
 Prepare the signals needed by admission before checking. A pending-prefill check
 uses per-engine uncached work when a prefix is known, and full input otherwise.
 Decode capacity uses the expected peak sequence length when available, including
 on a cache hit. Power-of-two retains the selected engine's load record from
 selection and passes it to admission without another snapshot. A single candidate
 still has its load read for admission, even though selection needs no comparison.
 Neither the bucket nor HTTP handler supplies observations. Concrete load-aware
 acceptance rules and additional cache-specific admission signals remain follow-up
 work. Synchronous checks do not fetch telemetry over the network themselves.
 `AllowAll` is the default for new explicit policy attachments. It leaves health,
 role, membership, and policy preferences in force. Migrated configurations must
 retain their existing capacity and configured budget checks; see compatibility
 below. Each check defines its missing-data behavior. Unknown load is not zero;
 the existing capacity and pending-prefill checks allow requests without a fresh,
 complete native report.
 The cache policy's `worker_queue_limit` is a **soft preference**, not
 `QueueLimitAdmission`. Saturation handling can reconsider a queued engine, but
 cannot bypass attached hard admission.
 Admission checks observe capacity; they do not reserve it. Concurrent requests
 may pass against the same observation. Strict reservations would require a
 separate mechanism.
 ## 5. Concrete policies
 | Policy | Selection behavior |
 | --- | --- |
 | `RoundRobinPolicy` | Rotate over candidates using a cursor owned by this policy instance |
 | `RandomPolicy` | Choose uniformly from candidates |
 | `PowerOfTwoPolicy` | Sample two distinct candidates when possible and choose the lower-pressure engine using the stage's load comparison |
 | `LeastLoadPolicy` (`load_based`) | Choose the least loaded engine; preserve tie-breaking, telemetry fallback, and recent-dispatch correction |
 | `SessionAwarePolicy` | Reuse an admitted session binding; use power-of-two for new or keyless sessions |
 | `StickyPolicy` | Reuse an admitted routing-key binding; use the configured fallback for new or missing keys |
 | `CacheAwarePolicy` | Prefer a usable prefix under cache and pressure rules; use a load-based fallback on a miss |
 Session and sticky policies do not create assignments for missing keys. A
 binding outside the candidates cannot win. A missing binding may invoke policy
 fallback within the group; hard admission rejection remains an error.
 Sticky fallback supports `round_robin`, `random`, `power_of_two`, and `load_based`,
 with round-robin as the default. Nested fallbacks use `AllowAll`; the owning
 policy explicitly checks the engine returned by its fallback.
 ### Cache-aware behavior
 The architecture preserves the cache algorithm within the supplied candidate
 set. Its responsibilities are:
 1. Look up prefix ownership through the local index or remote indexer.
 2. Apply minimum matched-token and optional ratio thresholds.
 3. Bound candidates using prefix/pressure ordering and the configured minimum,
   ratio, and maximum worker counts.
 4. Apply the soft queue gate and saturation rules, and call admission explicitly
   as required by the cache policy's candidate-selection algorithm.
 5. Choose among usable prefix holders using uncached work, the switch margin,
   and the pressure guard.
 6. On a miss, run the load fallback, preferring engines admitted by the soft
   queue gate when available.
 Candidate limits and saturation observations use only the selected bucket's
 role-group candidates.
 Saturation pinning must still pass hard admission.
 The target load fallback supports power-of-k sampling through
 `--min-load-choices`, default 2. When k covers the group, choose the exact minimum.
 Preserve queue-tier preference and avoid sorting with a pairwise pressure
 comparator that does not define a total ordering.
 Memoize the prefix lookup once per request, including remote I/O. Each policy
 restricts those matches to its own candidates. A memoized lookup does not imply
 that another bucket or stage's cache selection and admission have already run.
 Any optimization that skips those steps must establish that the previous result
 applies to the current group.
 ## 6. Shared state and construction
 Application wiring starts shared services once. Policy construction validates
 configuration and passes the required handles to each policy and admission
 implementation. For example, `PowerOfTwoPolicy::new(Arc<EngineReportedLoadTable>)`
 retains the application's shared load table. KV-aware and affinity-aware policies
 receive their corresponding shared handles when implemented. Policies with no
 state dependency require none. Policy instances do not create duplicate
 subscriptions, polling loops, indexes, or remote-client concurrency limits.
 Requirements come from all configured policies, their admission checks, and
 nested fallbacks. This includes tokenization, affinity-header extraction, load
 observations, and dispatch timestamps. A role-group override must receive its
 required settings even when the model's default uses a different policy.
 ### Load state
 `state/load_monitor/` owns engine reports and existing router-local request
 accounting. Power-of-two owns an `Arc<EngineReportedLoadTable>` and captures a snapshot
 locally for each nonempty selection attempt. Its `pick` method selects the engine,
 then passes that engine's borrowed load record directly to admission.
 No snapshot or observation is added to `Pick`, `PickRequest`,
 or the bucket interface, and no shared observation context is threaded through
 policies or fallbacks.
 The existing snapshot reader preserves rank aggregation, freshness, and capacity
 fields. Missing, stale, or rank-incomplete reports yield `None`, not zero load.
 A new pick reads current state. Admission reuses the selected observation even if
 reports change after selection; it neither recaptures nor reserves capacity.
 Fallback policies read their own state and do not share snapshots with callers.
 Snapshot capture still scans the full table; an engine-scoped reader can be added
 if profiling justifies it. Power-of-two reuses the legacy prefill/decode pressure
 comparisons, including router-local fallback. Concrete load-aware admission remains
 in #40271. Further shared load interpretation and correction for dispatches since
 the report remain follow-ups; these must preserve source, freshness, and available
 measurements without adding another
 independent in-flight counter. Load and cache observations are not an atomic global
 snapshot. Preserve request-guard cleanup.
 ### Cache state
 `state/kv_events/` owns local subscriptions, event application, hashing, and the
 radix-tree index. The tree stores prefix ownership and storage tiers, not KV
 tensors. Eviction, invalidation, and worker removal update this shared state.
 A small `CacheSource` adapter beside the cache policy normalizes local and remote
 lookups into common prefix results. Preserve model/cache namespaces, required
 rank information, and hash/block-size handling. Convert block counts to token
 counts only with a valid conversion.
 A confirmed miss and an unavailable backend both allow cache fallback, but
 remain distinguishable in diagnostics. Invalid queries and configuration errors
 remain explicit errors. Remote indexing and load-only deployments do not need
 a duplicate local cache tree.
 ### Affinity state
 `AffinityStore` owns scoped assignments, idle expiry, worker invalidation, and
 atomic updates. Policies decide when to reuse or replace an assignment.
 Concurrent first assignments must converge on an effective binding that is
 still a candidate and passes admission. If a concurrent binding returns a
 different engine, the policy must check that engine before returning it.
 Reconcile conflicts with bounded retry.
 Create or replace a binding only after admission succeeds. A binding records
 preferred placement, not successful execution, so it may remain if later PD
 selection or dispatch fails. It must not increment dispatch accounting.
 ## 7. Configuration and compatibility
 This conceptual example shows the target attachment model. It is not copyable
 current CLI/JSON syntax; existing bucket field names need not change.
 ```yaml
 buckets:
  - id: short-context
    rank: 10
    limits: {min: 0, max: 4096}
    max_context_tokens: 8192
    groups:
      pd:
        prefill:
          worker_ids: [P1, P2]
          policy:
            type: cache_aware
            admission: {type: capacity}
        decode:
          worker_ids: [D1, D2]
          policy:
            type: power_of_two
            admission: {type: capacity}
  - id: long-context
    rank: 20
    limits: {min: 0, max: 131072}
    max_context_tokens: 131072
    groups:
      pd:
        prefill:
          worker_ids: [P3, P4]
          policy:
            type: cache_aware
            admission: {type: capacity}
        decode:
          worker_ids: [D3, D4]
          policy:
            type: power_of_two
            admission: {type: capacity}
 ```
 A request with 4k input tokens and a 16k expected peak cannot fit the short
 bucket's context capacity. It selects the long bucket and both of its P/D groups.
 With a known peak of 8k or less, the same input selects both groups of the short bucket.
 The planned factory validates unique nonempty bucket IDs, token ranges, and
 role-compatible membership. Each bucket is either plain or PD. The selected
 bucket invokes its required groups through `pick_engines`. The existing worker registry
 still rejects mixed plain and PD engines within one model; this PR preserves
 that constraint. The engine group's model and stage filters apply on every pick.
 Legacy `BucketSpec` represents a single role-specific membership set. Migration
 must explicitly associate prefill and decode specs into complete PD buckets;
 never infer those associations from matching rank or similar names. Translation
 of role-specific ranges/ranks into bucket-level constraints needs explicit
 validation and is deferred with the configuration factory.
 Retained settings keep their meanings, defaults, units, and validation unless a
 change is listed below. Policy-specific tuning applies to role groups using that
 policy. Reject unsupported settings and incompatible combinations at startup;
 do not accept and ignore them.
 ### Retained behavior
 - Keep `load_based` as the CLI name for `LeastLoadPolicy`.
 - Preserve explicit engine membership and context constraints. Bucket-level
  ranges and ordering replace independent per-stage selection.
  Restore existing SLO behavior in the separate SLO PR before serving switchover;
  legacy routing continues to support SLOs during this skeleton-only phase.
 - Preserve cache-provider selection, endpoint validation, query timeout and
  concurrency limits, and unavailable-backend fallback.
 - Preserve cache thresholds and tuning: the 1,024-token default minimum hit,
  optional ratio gate, candidate bounds, switch margin, pressure guard, soft
  queue limit, and saturation floor.
 - Preserve session and sticky headers, idle timeouts, eviction cadence, and the
  four sticky fallback choices. Global modes need a bucket-first migration design.
 - Translate `--filter overloaded` and `--max-in-flight` into
  `InFlightLimitAdmission`, composed with other checks through `AllOfAdmission`.
 - Preserve configured capacity, pending-prefill, and in-flight checks, including
  their missing-report behavior. Power-of-two applies admission to its selected
  engine; other policies explicitly place checks in their selection logic.
  Other paths use `AllowAll` unless a check is configured. Never silently discard
  a configured budget.
 Listener and shutdown configuration, discovery, worker health and circuit
 breakers, tokenizer loading, request timeouts, sampling overrides, and logging
 remain outside the policy redesign. Preserve their existing behavior and
 validation, including dispatch-time breaker probes and request cancellation.
 ### Deliberate changes
 | Behavior | Target |
 | --- | --- |
 | Power-of-two admission | Check only the chosen engine; rejection advances to the next bucket |
 | Round-robin cursor | One cursor per role-group policy instance |
 | Capacity exhaustion | Try the next compatible bucket; return accumulated rejection details if all fail |
 | Primary/backup proposals and post-policy substitution | Removed; each policy returns one engine |
 | Session affinity | Reuse admitted bindings; remove primary/backup pressure escape |
 | Omitted `--affinity-mode` | Admitted-binding reuse replaces the former soft-mode default |
 | Pressure-guard tuning | Applies to cache-aware selection; reject session-only use |
 | Policy attachment | Explicit role-group policy overrides the applicable model/stage default |
 Reject these dropped options explicitly:
 - `--policy fused_score` and `--policy score_policy`, including `--fuse` terms
  and weights.
 - `--decode-policy legacy_host_affinity`.
 - `--stable-pair`.
 - `--affinity-mode soft`; only strict admitted-binding reuse remains.
 - `--filter prefix_cache` and `--prefix-cache-min-share`. The removed prefix-share
  filter is not equivalent to the cache-aware minimum-hit gate.
 ## 8. Dispatch and failure handling
 A successful pick means admission passed against the observed state. Health and
 capacity can change before dispatch. The handler owns network operations, retry
 rules, and accounting for the engines actually dispatched to.
 Selection fallback advances through the ordered compatible buckets before any
 network dispatch. Every PD attempt picks both engines from that bucket. Transport
 retry integration remains separate work; a forwarding failure does not resume
 the bucket loop or silently replace a successful policy pick.
 For PD, acquire and release accounting for the actual stages and clean up
 partial setup on failure. Policy selection does not own the PD request lifetime.
 Selection metrics must not imply that dispatch or execution succeeded.
 ## Implementation status
 This PR adds the side-by-side interfaces in `src/buckets_reorg.rs` and
 `src/policies_reorg/`, plus a configurable bucket-first implementation behind
 `chat_completions`. The live `src/policies/` path remains the default.
 Implemented here:
 - `BucketResolver::resolve` returns all length-compatible buckets in capacity/rank/ID order.
 - `Bucket::pick_engines` owns plain/PD orchestration and stage-specific policy
  requests; `BucketRequest` carries prepared facts and `BucketPick` retains picks.
 - `Bucket` owns input limits, context capacity, rank, and plain-or-PD groups.
 - `EngineGroup::pick` owns live candidate filtering, policy invocation, and
  exact candidate validation, without cross-bucket fallback.
 - `Policy::pick`, within-group fallback interface, per-engine `EngineAdmission::check`,
  and `AllowAll`. Power-of-two samples two distinct engines, compares stage pressure,
  and checks its selected engine with no replacement on rejection.
 - Policy-owned load dependency and local observations. Power-of-two passes the
  selected engine's load record directly to admission, without another snapshot.
  `PickRequest`, `Pick`, and bucket APIs carry no load observations.
 - The reorg chat implementation iterates resolved buckets, calls `pick_engines`,
  advances on empty candidates/admission rejection, and
  dispatches only after one complete selection. Exhaustion retains admission reasons.
 - `AppContext::chat_routing` configures legacy versus reorg routing on the same
  endpoint and carries the reorg model-resolver map.
 Follow-up order: concrete admission (#40271), then bucket SLO ordering
 in a separate PR, followed by remaining policies and production configuration.
 Not yet implemented in the reorg path:
 - Other concrete policies and capacity/in-flight admission checks.
 - SLO estimates, targets, and bucket preference ordering.
 - CLI/configuration parsing, validation, and model-specific construction.
  The YAML above is illustrative; reorg resolvers are installed in code.
 - Session modes, prefix memoization, and cache-aware selection.
 - Shared load interpretation, dispatch correction, and policy-specific
  dispatch-timestamp requirements.
 - PD compatibility filtering, retry integration, and legacy-route switchover.
 The preceding policy sections describe target behavior for those follow-ups;
 they do not claim those capabilities are present in this PR.
@@ -0,0 +1,238 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 //! Order buckets by request length; each bucket owns the groups used to pick engines.
 //!
 //! ```text
 //! BucketResolver (one model's buckets)
 //!   -> Bucket (token limits, context capacity, rank)
 //!        -> Plain: one EngineGroup
 //!        -> PD: prefill + decode EngineGroups
 //!             -> each EngineGroup: worker membership + its own Policy
 //! ```
 //!
 //! - [`BucketResolver::resolve`] returns length-compatible buckets in preference order.
 //! - [`EngineGroup::pick`] filters live workers by model, health, stage, and membership,
 //!   then calls [`Policy::pick`] and validates the returned engine.
 //!
 //! The handler tries buckets in order, advancing on missing candidates or admission
 //! rejection. Both P/D picks must succeed in the same bucket before dispatch.
 //! [`WorkerRegistry`] owns live workers; groups reference their IDs. Policies own
 //! their load/KV/affinity dependencies and share observations within each attempt.
 use std::collections::HashSet;
 use std::sync::Arc;
 use crate::discovery::{ModelId, WorkerId};
 use crate::policies_reorg::{Pick, PickError, PickRequest, Policy, Stage};
 use crate::workers::WorkerRegistry;
 #[derive(Debug, Clone, Copy, Default)]
 pub struct TokenLimits {
    pub min: Option<u64>,
    pub max: Option<u64>,
 }
 impl TokenLimits {
    fn fits(&self, tokens: u64) -> bool {
        self.min.is_none_or(|min| tokens >= min) && self.max.is_none_or(|max| tokens <= max)
    }
 }
 #[derive(Debug)]
 pub struct EngineGroup {
    /// `None` includes all registered engines matching the request's model and role.
    pub worker_ids: Option<HashSet<WorkerId>>,
    pub policy: Arc<dyn Policy>,
 }
 impl EngineGroup {
    pub fn new(policy: Arc<dyn Policy>) -> Self {
        Self {
            worker_ids: None,
            policy,
        }
    }
    /// Resolve live members and invoke this group's policy. Never changes buckets.
    pub async fn pick(
        &self,
        workers: &WorkerRegistry,
        request: &PickRequest<'_>,
    ) -> Result<Pick, PickError> {
        let mut engines: Vec<_> = workers
            .healthy_workers_for(request.model)
            .into_iter()
            .filter(|engine| engine.mode() == request.stage)
            .filter(|engine| {
                self.worker_ids
                    .as_ref()
                    .is_none_or(|ids| ids.contains(&engine.id))
            })
            .collect();
        // Stable order so cursor-based policies see a consistent candidate list.
        engines.sort_by(|left, right| left.id.0.cmp(&right.id.0));
        if engines.is_empty() {
            return Err(PickError::NoCandidates);
        }
        let pick = self.policy.pick(&engines, request).await?;
        if !engines
            .iter()
            .any(|engine| Arc::ptr_eq(engine, &pick.engine))
        {
            return Err(PickError::OutsideCandidates(pick.engine.id.clone()));
        }
        Ok(pick)
    }
 }
 /// A bucket serves a request on one plain engine or on its own P/D groups.
 #[derive(Debug)]
 pub enum BucketGroups {
    Plain(EngineGroup),
    Pd {
        prefill: EngineGroup,
        decode: EngineGroup,
    },
 }
 /// Prepared request facts shared by all bucket attempts. The bucket supplies
 /// its ID and each group's stage when calling policies.
 #[derive(Debug)]
 pub struct BucketRequest<'a> {
    pub model: &'a ModelId,
    pub input_tokens: u64,
    pub expected_peak_tokens: Option<u64>,
    pub token_ids: Option<&'a [u32]>,
    pub session_key: Option<&'a str>,
    pub routing_key: Option<&'a str>,
 }
 /// A complete selection from one bucket. For plain serving, `prefill` is the
 /// plain engine and `decode` is absent; PD supplies both picks.
 #[derive(Debug)]
 pub struct BucketPick {
    pub prefill: Pick,
    pub decode: Option<Pick>,
 }
 #[derive(Debug)]
 pub struct Bucket {
    pub id: String,
    /// Break ties between equally sized buckets; lower ranks win.
    pub rank: u32,
    /// Inclusive input-token range used to choose the bucket.
    pub limits: TokenLimits,
    /// Full sequence capacity, checked against the expected peak when known.
    pub max_context_tokens: Option<u64>,
    pub groups: BucketGroups,
 }
 impl Bucket {
    pub fn new(id: impl Into<String>, groups: BucketGroups) -> Self {
        Self {
            id: id.into(),
            rank: 0,
            limits: TokenLimits::default(),
            max_context_tokens: None,
            groups,
        }
    }
    /// Select this bucket's plain engine or complete P/D pair, without dispatching.
    /// A failed group reports its stage; the caller may then try another bucket.
    pub async fn pick_engines(
        &self,
        workers: &WorkerRegistry,
        request: &BucketRequest<'_>,
    ) -> Result<BucketPick, (Stage, PickError)> {
        let (prefill, decode) = match &self.groups {
            BucketGroups::Plain(group) => (
                self.pick_from_group(group, Stage::Plain, workers, request)
                    .await?,
                None,
            ),
            BucketGroups::Pd { prefill, decode } => {
                let prefill = self
                    .pick_from_group(prefill, Stage::Prefill, workers, request)
                    .await?;
                let decode = self
                    .pick_from_group(decode, Stage::Decode, workers, request)
                    .await?;
                (prefill, Some(decode))
            }
        };
        Ok(BucketPick { prefill, decode })
    }
    /// Scope the request to this bucket and role, then ask the group for one engine.
    async fn pick_from_group(
        &self,
        group: &EngineGroup,
        stage: Stage,
        workers: &WorkerRegistry,
        request: &BucketRequest<'_>,
    ) -> Result<Pick, (Stage, PickError)> {
        let request = PickRequest {
            model: request.model,
            stage,
            bucket: &self.id,
            input_tokens: request.input_tokens,
            expected_peak_tokens: request.expected_peak_tokens,
            token_ids: request.token_ids,
            session_key: request.session_key,
            routing_key: request.routing_key,
        };
        group
            .pick(workers, &request)
            .await
            .map_err(|error| (stage, error))
    }
    fn fits(&self, input_tokens: u64, expected_peak_tokens: Option<u64>) -> bool {
        self.limits.fits(input_tokens)
            && self
                .max_context_tokens
                .is_none_or(|max| expected_peak_tokens.unwrap_or(input_tokens) <= max)
    }
    fn input_capacity(&self) -> u64 {
        self.limits
            .max
            .unwrap_or(u64::MAX)
            .min(self.max_context_tokens.unwrap_or(u64::MAX))
    }
 }
 /// Model-specific bucket configuration. Selection does not inspect engine state.
 #[derive(Debug, Default)]
 pub struct BucketResolver {
    pub buckets: Vec<Bucket>,
 }
 impl BucketResolver {
    pub fn new(buckets: Vec<Bucket>) -> Self {
        Self { buckets }
    }
    /// Return all length-compatible buckets, ordered by input capacity, rank, and ID.
    /// The caller tries their groups in order until a complete engine selection succeeds.
    pub fn resolve(
        &self,
        input_tokens: u64,
        expected_peak_tokens: Option<u64>,
    ) -> Result<Vec<&Bucket>, PickError> {
        if expected_peak_tokens.is_some_and(|tokens| tokens < input_tokens) {
            return Err(PickError::InvalidSignal(
                "expected peak tokens are below input length".into(),
            ));
        }
        let mut buckets: Vec<_> = self
            .buckets
            .iter()
            .filter(|bucket| bucket.fits(input_tokens, expected_peak_tokens))
            .collect();
        buckets.sort_by_key(|bucket| (bucket.input_capacity(), bucket.rank, &bucket.id));
        Ok(buckets)
    }
 }
@@ -8,10 +8,12 @@
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 pub mod buckets_reorg;
 pub mod config;
 pub mod discovery;
 pub mod health;
 pub mod policies;
 pub mod policies_reorg;
 pub mod proxy;
 pub mod server;
 pub mod state;
@@ -0,0 +1,42 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use std::fmt::Debug;
 use crate::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
 use crate::workers::Worker;
 use super::{PickError, PickRequest};
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Decision {
    Allow,
    Reject(String),
 }
 /// Checks one engine using the load observation retained by selection.
 /// `None` means no usable load observation, never zero load. Each check defines
 /// its missing-data behavior and owns any other state handles it needs.
 /// Each policy decides when to check an engine and how to handle rejection.
 pub trait EngineAdmission: Send + Sync + Debug {
    fn check(
        &self,
        engine: &Worker,
        request: &PickRequest<'_>,
        load: Option<&EngineReportedWorkerLoad>,
    ) -> Result<Decision, PickError>;
 }
 #[derive(Debug)]
 pub struct AllowAll;
 impl EngineAdmission for AllowAll {
    fn check(
        &self,
        _: &Worker,
        _: &PickRequest<'_>,
        _: Option<&EngineReportedWorkerLoad>,
    ) -> Result<Decision, PickError> {
        Ok(Decision::Allow)
    }
 }
@@ -0,0 +1,105 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 //! Side-by-side implementation of POLICY_DESIGN.md. Chat routing can opt into
 //! this interface through AppContext; `policies` remains the default.
 pub mod admission;
 pub mod power_of_two;
 use std::fmt::Debug;
 use std::sync::Arc;
 use futures::future::BoxFuture;
 use crate::discovery::{ModelId, WorkerId};
 use crate::workers::Worker;
 pub use crate::discovery::WorkerMode as Stage;
 /// Request facts for engine selection. Policies own their shared-state handles.
 #[derive(Debug, Clone, Copy)]
 pub struct PickRequest<'a> {
    pub model: &'a ModelId,
    pub stage: Stage,
    pub bucket: &'a str,
    pub input_tokens: u64,
    pub expected_peak_tokens: Option<u64>,
    pub token_ids: Option<&'a [u32]>,
    pub session_key: Option<&'a str>,
    pub routing_key: Option<&'a str>,
 }
 impl<'a> PickRequest<'a> {
    pub fn new(model: &'a ModelId, stage: Stage, input_tokens: u64) -> Self {
        Self {
            model,
            stage,
            bucket: "",
            input_tokens,
            expected_peak_tokens: None,
            token_ids: None,
            session_key: None,
            routing_key: None,
        }
    }
 }
 #[derive(Debug, Clone)]
 pub struct Pick {
    pub engine: Arc<Worker>,
    pub reason: &'static str,
 }
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Rejection {
    pub engine: WorkerId,
    pub reason: String,
 }
 #[derive(Debug, thiserror::Error)]
 pub enum PickError {
    #[error("no bucket matches the request length")]
    NoMatchingBucket,
    #[error("no candidates")]
    NoCandidates,
    #[error("no admissible engine: {0:?}")]
    NoAdmissibleEngine(Vec<Rejection>),
    #[error("selected engine rejected: {0:?}")]
    AdmissionRejected(Rejection),
    #[error("invalid signal: {0}")]
    InvalidSignal(String),
    #[error("invalid configuration: {0}")]
    InvalidConfiguration(String),
    #[error("policy selected an engine outside its candidates: {0:?}")]
    OutsideCandidates(WorkerId),
 }
 /// Returns one admitted engine from exactly the supplied candidates.
 /// Implementations receive shared load, KV, and affinity handles at construction;
 /// they obtain their own observations rather than asking callers to supply them.
 pub trait Policy: Send + Sync + Debug {
    /// Read required state locally and pass the selected engine's observations to admission.
    fn pick<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>>;
    /// Runs on a miss within the same candidates; never on an admission rejection.
    fn fallback(&self) -> Option<&dyn Policy> {
        None
    }
    /// Delegate within the same candidates; the fallback reads its own state.
    fn pick_fallback<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>> {
        match self.fallback() {
            Some(fallback) => fallback.pick(engines, request),
            None => Box::pin(async { Err(PickError::NoCandidates) }),
        }
    }
 }
@@ -0,0 +1,79 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use std::sync::Arc;
 use std::time::Instant;
 use futures::future::BoxFuture;
 use rand::Rng;
 use crate::policies::admission::{compare_decode_pressure, compare_prefill_pressure};
 use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
 use crate::workers::Worker;
 use super::admission::{AllowAll, Decision, EngineAdmission};
 use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
 /// Samples two distinct engines and selects the one with lower stage pressure.
 /// Checks admission only on the selected engine; rejection never resamples.
 #[derive(Debug)]
 pub struct PowerOfTwoPolicy {
    /// Shared application state; snapshots are local to each pick.
    engine_load: Arc<EngineReportedLoadTable>,
    pub admission: Arc<dyn EngineAdmission>,
 }
 impl PowerOfTwoPolicy {
    pub fn new(engine_load: Arc<EngineReportedLoadTable>) -> Self {
        Self {
            engine_load,
            admission: Arc::new(AllowAll),
        }
    }
 }
 impl Policy for PowerOfTwoPolicy {
    fn pick<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>> {
        Box::pin(async move {
            if engines.is_empty() {
                return Err(PickError::NoCandidates);
            }
            // Selection and admission use the same load observation.
            let load = self.engine_load.capture_snapshot(Instant::now());
            let engine = match engines {
                [engine] => Arc::clone(engine),
                _ => {
                    let mut rng = rand::thread_rng();
                    let i = rng.gen_range(0..engines.len());
                    let mut j = rng.gen_range(0..engines.len() - 1);
                    if j >= i {
                        j += 1;
                    }
                    let (left, right) = (&engines[i], &engines[j]);
                    let pressure = match request.stage {
                        Stage::Plain | Stage::Prefill => {
                            compare_prefill_pressure(left, right, Some(&load))
                        }
                        Stage::Decode => compare_decode_pressure(left, right, Some(&load)),
                    };
                    Arc::clone(if pressure.is_gt() { right } else { left })
                }
            };
            let engine_load = load.fresh_load_for_url(&engine.url);
            if let Decision::Reject(reason) = self.admission.check(&engine, request, engine_load)? {
                return Err(PickError::AdmissionRejected(Rejection {
                    engine: engine.id.clone(),
                    reason,
                }));
            }
            Ok(Pick {
                engine,
                reason: "power_of_two",
            })
        })
    }
 }
@@ -1,7 +1,9 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use crate::buckets_reorg::BucketResolver;
 use crate::config::Config;
 use crate::discovery::ModelId;
 use crate::policies::buckets::BucketSelector;
 use crate::policies::prefix_provider::RadixTreePrefixProvider;
@@ -14,6 +16,7 @@ use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
 use crate::state::load_monitor::router_inflight_load::RouterInflightLoadRegistry;
 use crate::tokenizer::TokenizerRegistry;
 use crate::workers::WorkerRegistry;
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
@@ -26,6 +29,15 @@ const READINESS_NOT_READY: u8 = 0;
 const READINESS_READY: u8 = 1;
 const READINESS_DRAINING: u8 = 2;
 /// Routing implementation used by the standard chat-completions endpoint.
 /// Reorg configuration is installed explicitly until its CLI factory is available.
 #[derive(Debug, Default)]
 pub enum ChatRouting {
    #[default]
    Legacy,
    Reorg(HashMap<ModelId, BucketResolver>),
 }
 pub struct AppContext {
    pub config: Config,
    pub tokenizers: Arc<TokenizerRegistry>,
@@ -34,6 +46,8 @@ pub struct AppContext {
    pub policies: Arc<PolicyRegistry>,
    /// Converts static Bucket configuration into request candidate domains.
    pub bucket_selector: Arc<BucketSelector>,
    /// Select legacy policies or model-specific bucket-first routing.
    pub chat_routing: ChatRouting,
    /// Per-worker active-load bookkeeping shared by the proxy, policies,
    /// timeout janitor, and metrics.
    pub router_inflight_load: Arc<RouterInflightLoadRegistry>,
@@ -106,6 +120,7 @@ impl AppContext {
            registry,
            policies,
            bucket_selector,
            chat_routing: ChatRouting::Legacy,
            router_inflight_load,
            metrics,
            prefix_index: None,
@@ -195,6 +210,7 @@ impl AppContext {
            registry: Arc::new(WorkerRegistry::default()),
            policies: Arc::new(PolicyRegistry::default()),
            bucket_selector: Arc::new(BucketSelector::new(None)),
            chat_routing: ChatRouting::Legacy,
            router_inflight_load: RouterInflightLoadRegistry::with_defaults(),
            metrics: MetricsRegistry::new(),
            prefix_index: None,
@@ -3,6 +3,7 @@
 mod forward;
 mod preparation;
 mod reorg;
 use crate::config::{SessionAffinityMode, DEFAULT_MIN_LOAD_CHOICES};
 use crate::discovery::{ModelId, WorkerMode};
@@ -11,7 +12,7 @@ use crate::policies::selection::{
    select_decode_peer, select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
 };
 use crate::policies::{ExternalPrefixSignal, Policy};
-use crate::server::app_context::AppContext;
+use crate::server::app_context::{AppContext, ChatRouting};
 use crate::server::error::ApiError;
 use crate::server::metrics::PolicySelectionFailureReason;
 use crate::state::kv_events::{compute_block_hashes, compute_block_hashes_bigram};
@@ -38,6 +39,19 @@ pub async fn chat_completions(
    State(ctx): State<Arc<AppContext>>,
    headers: HeaderMap,
    body: Bytes,
 ) -> Result<Response<Body>, ApiError> {
    match &ctx.chat_routing {
        ChatRouting::Legacy => chat_completions_legacy(&ctx, headers, body).await,
        ChatRouting::Reorg(resolvers) => {
            reorg::chat_completions(&ctx, resolvers, headers, body).await
        }
    }
 }
 async fn chat_completions_legacy(
    ctx: &AppContext,
    headers: HeaderMap,
    body: Bytes,
 ) -> Result<Response<Body>, ApiError> {
    let start = Instant::now();
    let mut fields = parse_routing_fields(&body)?;
@@ -59,11 +73,11 @@ pub async fn chat_completions(
        .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?;
    let request =
-        PreparedChatRequest::prepare(&ctx, model, fields, body, policy.needs_request_tokens())?;
+        PreparedChatRequest::prepare(ctx, model, fields, body, policy.needs_request_tokens())?;
    // Pick a plain worker, or a prefill worker followed by a decode peer in PD mode.
    let workers = select_workers(
-        &ctx,
+        ctx,
        &request,
        &headers,
        policy.as_ref(),
@@ -73,7 +87,7 @@ pub async fn chat_completions(
    .await?;
    // PD sends to both workers and returns the decode response.
-    forward_chat_request(&ctx, request, workers, headers, start).await
+    forward_chat_request(ctx, request, workers, headers, start).await
 }
 fn pool_error(error: PdResolveError, model: &ModelId) -> ApiError {
@@ -0,0 +1,136 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use super::forward::{forward_chat_request, SelectedWorkers};
 use super::nonempty_header;
 use super::preparation::{parse_routing_fields, PreparedChatRequest};
 use crate::buckets_reorg::{BucketRequest, BucketResolver};
 use crate::discovery::ModelId;
 use crate::policies_reorg::{PickError, Stage};
 use crate::server::app_context::AppContext;
 use crate::server::error::ApiError;
 use axum::body::Body;
 use axum::http::{HeaderMap, Response};
 use bytes::Bytes;
 use std::collections::HashMap;
 use std::time::Instant;
 /// Bucket-first implementation selected by `AppContext::chat_routing`.
 pub(super) async fn chat_completions(
    ctx: &AppContext,
    resolvers: &HashMap<ModelId, BucketResolver>,
    headers: HeaderMap,
    body: Bytes,
 ) -> Result<Response<Body>, ApiError> {
    let start = Instant::now();
    let mut fields = parse_routing_fields(&body)?;
    let model = ModelId(
        fields
            .model
            .take()
            .ok_or_else(|| ApiError::BadRequest("missing `model` field".into()))?,
    );
    let resolver = resolvers
        .get(&model)
        .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?;
    // Length-based routing needs tokenization even for load-only group policies.
    let request = PreparedChatRequest::prepare(ctx, model, fields, body, true)?;
    let input_tokens = request.input_token_count as u64;
    let expected_peak_tokens = request
        .max_output_tokens
        .map(|output| {
            input_tokens.checked_add(output).ok_or_else(|| {
                ApiError::BadRequest("input and output token counts overflow".into())
            })
        })
        .transpose()?;
    let buckets = resolver
        .resolve(input_tokens, expected_peak_tokens)
        .map_err(|error| selection_error(error, &request.model, None))?;
    if buckets.is_empty() {
        return Err(selection_error(
            PickError::NoMatchingBucket,
            &request.model,
            None,
        ));
    }
    let bucket_request = BucketRequest {
        model: &request.model,
        input_tokens,
        expected_peak_tokens,
        token_ids: request.tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
        session_key: ctx
            .config
            .model
            .affinity
            .as_ref()
            .and_then(|config| nonempty_header(&headers, &config.session_id_header)),
        routing_key: ctx
            .config
            .model
            .sticky
            .as_ref()
            .and_then(|config| nonempty_header(&headers, &config.header_name)),
    };
    let mut rejections: Option<Vec<_>> = None;
    let mut missing_stage = None;
    for bucket in buckets {
        match bucket.pick_engines(&ctx.registry, &bucket_request).await {
            Ok(picks) => {
                // Dispatch only after this bucket supplies the entire plain or PD selection.
                let workers = SelectedWorkers {
                    prefill: picks.prefill.engine,
                    decode: picks.decode.map(|pick| pick.engine),
                    track_dispatch_timestamps: false,
                };
                return forward_chat_request(ctx, request, workers, headers, start).await;
            }
            Err((stage, error)) => {
                tracing::debug!(bucket = %bucket.id, ?stage, %error, "bucket selection failed");
                match error {
                    PickError::NoCandidates => missing_stage = Some(stage),
                    PickError::NoAdmissibleEngine(reasons) => {
                        rejections.get_or_insert_with(Vec::new).extend(reasons);
                    }
                    PickError::AdmissionRejected(reason) => {
                        rejections.get_or_insert_with(Vec::new).push(reason);
                    }
                    error => return Err(selection_error(error, &request.model, Some(stage))),
                }
            }
        }
    }
    // Preserve admission exhaustion even if a later bucket has no candidates.
    let error = match rejections {
        Some(reasons) => PickError::NoAdmissibleEngine(reasons),
        None => PickError::NoCandidates,
    };
    Err(selection_error(error, &request.model, missing_stage))
 }
 fn selection_error(error: PickError, model: &ModelId, stage: Option<Stage>) -> ApiError {
    tracing::warn!(%model, ?stage, %error, "reorg selection failed");
    match error {
        PickError::NoMatchingBucket => {
            ApiError::BadRequest("no bucket supports the requested token length".into())
        }
        PickError::NoCandidates => match stage {
            Some(Stage::Prefill) => ApiError::NoPrefillWorkersAvailable {
                model: model.0.clone(),
            },
            Some(Stage::Decode) => ApiError::NoDecodeWorkersAvailable {
                model: model.0.clone(),
            },
            _ => ApiError::NoHealthyWorkers {
                model: model.0.clone(),
            },
        },
        PickError::NoAdmissibleEngine(_) | PickError::AdmissionRejected(_) => {
            ApiError::PolicySelectionFailed {
                model: model.0.clone(),
            }
        }
        error => ApiError::Internal(error.into()),
    }
 }
@@ -10,5 +10,8 @@
 mod discovery;
 mod health;
 mod policies;
 mod policies_reorg;
 mod policies_reorg_load;
 mod policies_reorg_power_of_two;
 mod tokenizer;
 mod workers;
@@ -0,0 +1,486 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use std::sync::{Arc, Mutex};
 use futures::future::BoxFuture;
 use sgl_router::buckets_reorg::{
    Bucket, BucketGroups, BucketRequest, BucketResolver, EngineGroup, TokenLimits,
 };
 use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
 use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission};
 use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
 use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
 use sgl_router::workers::{Worker, WorkerRegistry};
 #[derive(Debug)]
 struct TestPolicy {
    admission: Arc<dyn EngineAdmission>,
    result: Option<Arc<Worker>>,
    miss: bool,
    invalid: bool,
    calls: Mutex<Vec<String>>,
 }
 impl Default for TestPolicy {
    fn default() -> Self {
        Self {
            admission: Arc::new(AllowAll),
            result: None,
            miss: false,
            invalid: false,
            calls: Mutex::new(Vec::new()),
        }
    }
 }
 impl Policy for TestPolicy {
    fn pick<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>> {
        Box::pin(async move {
            self.calls.lock().unwrap().push(request.bucket.to_owned());
            if self.invalid {
                return Err(PickError::InvalidSignal("test signal".into()));
            }
            if self.miss {
                return Err(PickError::NoCandidates);
            }
            if engines.is_empty() {
                return Err(PickError::NoCandidates);
            }
            let engine = self.result.clone().unwrap_or_else(|| engines[0].clone());
            if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? {
                return Err(PickError::AdmissionRejected(Rejection {
                    engine: engine.id.clone(),
                    reason,
                }));
            }
            Ok(Pick {
                engine,
                reason: "test",
            })
        })
    }
 }
 #[derive(Debug)]
 struct Reject(&'static str);
 impl EngineAdmission for Reject {
    fn check(
        &self,
        engine: &Worker,
        _: &PickRequest<'_>,
        _: Option<&EngineReportedWorkerLoad>,
    ) -> Result<Decision, PickError> {
        Ok(if engine.id.0 == self.0 {
            Decision::Reject("full".into())
        } else {
            Decision::Allow
        })
    }
 }
 fn spec(id: &str, mode: Stage, model: &str) -> WorkerSpec {
    WorkerSpec {
        id: WorkerId(id.into()),
        url: format!("http://{id}"),
        mode,
        model_ids: vec![ModelId(model.into())],
        bootstrap_port: None,
    }
 }
 fn registry() -> Arc<WorkerRegistry> {
    let workers = Arc::new(WorkerRegistry::default());
    for (id, mode, model) in [
        ("a", Stage::Plain, "m"),
        ("b", Stage::Plain, "m"),
        ("unhealthy", Stage::Plain, "m"),
        ("other", Stage::Plain, "other"),
        ("p", Stage::Prefill, "pd"),
        ("d", Stage::Decode, "pd"),
    ] {
        workers.add(spec(id, mode, model)).unwrap();
    }
    let unhealthy = workers.get(&WorkerId("unhealthy".into())).unwrap();
    for _ in 0..3 {
        unhealthy.breaker.record_failure();
    }
    workers
 }
 fn group(members: &[&str], policy: Arc<dyn Policy>) -> EngineGroup {
    EngineGroup {
        worker_ids: Some(members.iter().map(|id| WorkerId((*id).into())).collect()),
        policy,
    }
 }
 fn bucket(id: &str, max: Option<u64>, policy: Arc<dyn Policy>) -> Bucket {
    let mut bucket = Bucket::new(id, BucketGroups::Plain(EngineGroup::new(policy)));
    bucket.limits.max = max;
    bucket
 }
 #[tokio::test]
 async fn groups_isolate_model_health_stage_and_membership() {
    let workers = registry();
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let group = EngineGroup::new(Arc::new(TestPolicy::default()));
    assert_eq!(
        group.pick(&workers, &request).await.unwrap().engine.id.0,
        "a"
    );
    workers.remove(&WorkerId("a".into()));
    assert_eq!(
        group.pick(&workers, &request).await.unwrap().engine.id.0,
        "b"
    );
    workers.remove(&WorkerId("b".into()));
    assert!(matches!(
        group.pick(&workers, &request).await,
        Err(PickError::NoCandidates)
    ));
    let pd = ModelId("pd".into());
    for (stage, expected) in [(Stage::Prefill, "p"), (Stage::Decode, "d")] {
        let request = PickRequest::new(&pd, stage, 10);
        let group = self::group(
            &["p", "d", "other", "unhealthy"],
            Arc::new(TestPolicy::default()),
        );
        assert_eq!(
            group.pick(&workers, &request).await.unwrap().engine.id.0,
            expected
        );
    }
    let empty = self::group(&[], Arc::new(TestPolicy::default()));
    assert!(matches!(
        empty.pick(&workers, &request).await,
        Err(PickError::NoCandidates)
    ));
 }
 #[test]
 fn resolve_orders_all_length_fits_by_capacity_rank_and_id() {
    let policy = Arc::new(TestPolicy::default());
    let mut z = bucket("z", Some(20), policy.clone());
    z.rank = 1;
    let mut a = bucket("a", Some(20), policy.clone());
    a.rank = 1;
    let mut later = bucket("later", Some(20), policy.clone());
    later.rank = 2;
    let mut min = bucket("min", Some(15), policy.clone());
    min.limits.min = Some(11);
    let resolver = BucketResolver::new(vec![
        bucket("catch-all", None, policy.clone()),
        bucket("too-small", Some(9), policy),
        z,
        later,
        a,
        min,
    ]);
    assert_eq!(
        resolver
            .resolve(10, None)
            .unwrap()
            .iter()
            .map(|bucket| bucket.id.as_str())
            .collect::<Vec<_>>(),
        ["a", "z", "later", "catch-all"]
    );
    assert_eq!(resolver.resolve(11, None).unwrap()[0].id, "min");
    assert_eq!(resolver.resolve(15, None).unwrap()[0].id, "min");
    assert_eq!(resolver.resolve(20, None).unwrap()[0].id, "a");
    assert_eq!(resolver.resolve(21, None).unwrap()[0].id, "catch-all");
 }
 #[test]
 fn context_capacity_checks_peak_when_known_and_input_otherwise() {
    let policy = Arc::new(TestPolicy::default());
    let mut short = bucket("short", None, policy.clone());
    short.max_context_tokens = Some(20);
    let mut long = bucket("long", None, policy);
    long.max_context_tokens = Some(30);
    let resolver = BucketResolver::new(vec![long, short]);
    assert_eq!(resolver.resolve(10, None).unwrap()[0].id, "short");
    assert_eq!(resolver.resolve(10, Some(20)).unwrap()[0].id, "short");
    assert_eq!(resolver.resolve(10, Some(21)).unwrap()[0].id, "long");
    assert!(resolver.resolve(10, Some(31)).unwrap().is_empty());
    assert!(resolver.resolve(31, None).unwrap().is_empty());
    assert!(matches!(
        resolver.resolve(10, Some(9)),
        Err(PickError::InvalidSignal(_))
    ));
    assert!(BucketResolver::default()
        .resolve(1, None)
        .unwrap()
        .is_empty());
 }
 #[tokio::test]
 async fn selected_pd_bucket_owns_both_memberships_and_policies() {
    let workers = registry();
    workers.add(spec("p2", Stage::Prefill, "pd")).unwrap();
    workers.add(spec("d2", Stage::Decode, "pd")).unwrap();
    let model = ModelId("pd".into());
    let prefill_policy = Arc::new(TestPolicy::default());
    let decode_policy = Arc::new(TestPolicy::default());
    let resolver = BucketResolver::new(vec![Bucket::new(
        "shared",
        BucketGroups::Pd {
            prefill: group(&["p2", "d", "a"], prefill_policy.clone()),
            decode: group(&["d2", "p", "other"], decode_policy.clone()),
        },
    )]);
    let bucket = resolver.resolve(10, Some(20)).unwrap()[0];
    let request = BucketRequest {
        model: &model,
        input_tokens: 10,
        expected_peak_tokens: Some(20),
        token_ids: None,
        session_key: None,
        routing_key: None,
    };
    let picks = bucket.pick_engines(&workers, &request).await.unwrap();
    assert_eq!(picks.prefill.engine.id.0, "p2");
    assert_eq!(picks.decode.unwrap().engine.id.0, "d2");
    assert_eq!(*prefill_policy.calls.lock().unwrap(), ["shared"]);
    assert_eq!(*decode_policy.calls.lock().unwrap(), ["shared"]);
 }
 #[tokio::test]
 async fn resolver_includes_empty_groups_without_invoking_policies() {
    let workers = registry();
    let model = ModelId("m".into());
    let policy = Arc::new(TestPolicy::default());
    let mut empty = Bucket::new(
        "empty",
        BucketGroups::Plain(group(&["missing"], policy.clone())),
    );
    empty.limits = TokenLimits {
        min: None,
        max: Some(10),
    };
    let resolver = BucketResolver::new(vec![empty, bucket("available", Some(20), policy.clone())]);
    let buckets = resolver.resolve(10, None).unwrap();
    assert_eq!(
        buckets
            .iter()
            .map(|bucket| bucket.id.as_str())
            .collect::<Vec<_>>(),
        ["empty", "available"]
    );
    let bucket = buckets[0];
    let BucketGroups::Plain(group) = &bucket.groups else {
        panic!("expected plain")
    };
    let request = PickRequest {
        bucket: &bucket.id,
        ..PickRequest::new(&model, Stage::Plain, 10)
    };
    assert!(matches!(
        group.pick(&workers, &request).await,
        Err(PickError::NoCandidates)
    ));
    assert!(policy.calls.lock().unwrap().is_empty());
 }
 #[tokio::test]
 async fn group_propagates_rejections_misses_and_invalid_signals() {
    let workers = registry();
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let rejected = group(
        &["a"],
        Arc::new(TestPolicy {
            admission: Arc::new(Reject("a")),
            ..Default::default()
        }),
    );
    assert!(matches!(
        rejected.pick(&workers, &request).await,
        Err(PickError::AdmissionRejected(_))
    ));
    let miss = group(
        &["a"],
        Arc::new(TestPolicy {
            miss: true,
            ..Default::default()
        }),
    );
    assert!(matches!(
        miss.pick(&workers, &request).await,
        Err(PickError::NoCandidates)
    ));
    let invalid = group(
        &["a"],
        Arc::new(TestPolicy {
            invalid: true,
            ..Default::default()
        }),
    );
    assert!(matches!(
        invalid.pick(&workers, &request).await,
        Err(PickError::InvalidSignal(_))
    ));
 }
 #[tokio::test]
 async fn group_rejects_foreign_pick_even_with_same_worker_id() {
    let workers = registry();
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let group = group(
        &["a"],
        Arc::new(TestPolicy {
            result: Some(Arc::new(Worker::new(spec("a", Stage::Plain, "m")))),
            ..Default::default()
        }),
    );
    assert!(matches!(
        group.pick(&workers, &request).await,
        Err(PickError::OutsideCandidates(_))
    ));
 }
 #[tokio::test]
 async fn selected_engine_rejection_does_not_try_an_alternative() {
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let engines = [
        Arc::new(Worker::new(spec("a", Stage::Plain, "m"))),
        Arc::new(Worker::new(spec("b", Stage::Plain, "m"))),
    ];
    let policy = TestPolicy {
        admission: Arc::new(Reject("a")),
        ..Default::default()
    };
    assert!(matches!(
        policy.pick(&engines, &request).await,
        Err(PickError::AdmissionRejected(_))
    ));
    assert!(matches!(
        policy.pick(&[], &request).await,
        Err(PickError::NoCandidates)
    ));
 }
 #[tokio::test]
 async fn bucket_scopes_plain_pick_and_preserves_request_facts() {
    #[derive(Debug)]
    struct InspectRequest;
    impl Policy for InspectRequest {
        fn pick<'a>(
            &'a self,
            engines: &'a [Arc<Worker>],
            request: &'a PickRequest<'a>,
        ) -> BoxFuture<'a, Result<Pick, PickError>> {
            Box::pin(async move {
                assert_eq!(request.model.0, "m");
                assert_eq!(request.bucket, "plain-bucket");
                assert_eq!(request.stage, Stage::Plain);
                assert_eq!(request.input_tokens, 2);
                assert_eq!(request.expected_peak_tokens, Some(12));
                assert_eq!(request.token_ids, Some([7, 9].as_slice()));
                assert_eq!(request.session_key, Some("session"));
                assert_eq!(request.routing_key, Some("routing"));
                Ok(Pick {
                    engine: engines[0].clone(),
                    reason: "inspected",
                })
            })
        }
    }
    let workers = registry();
    let model = ModelId("m".into());
    let bucket = Bucket::new(
        "plain-bucket",
        BucketGroups::Plain(group(&["b"], Arc::new(InspectRequest))),
    );
    let request = BucketRequest {
        model: &model,
        input_tokens: 2,
        expected_peak_tokens: Some(12),
        token_ids: Some(&[7, 9]),
        session_key: Some("session"),
        routing_key: Some("routing"),
    };
    let picks = bucket.pick_engines(&workers, &request).await.unwrap();
    assert_eq!(picks.prefill.engine.id.0, "b");
    assert_eq!(picks.prefill.reason, "inspected");
    assert!(picks.decode.is_none());
 }
 #[tokio::test]
 async fn power_of_two_checks_selected_engine_and_propagates_rejection_without_fallback() {
    use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
    use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
    #[derive(Debug)]
    struct Check {
        calls: Mutex<Vec<WorkerId>>,
        reject: bool,
        invalid: bool,
    }
    impl EngineAdmission for Check {
        fn check(
            &self,
            engine: &Worker,
            _: &PickRequest<'_>,
            _: Option<&EngineReportedWorkerLoad>,
        ) -> Result<Decision, PickError> {
            self.calls.lock().unwrap().push(engine.id.clone());
            if self.invalid {
                Err(PickError::InvalidSignal("admission input".into()))
            } else if self.reject {
                Ok(Decision::Reject("full".into()))
            } else {
                Ok(Decision::Allow)
            }
        }
    }
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let engine = Arc::new(Worker::new(spec("a", Stage::Plain, "m")));
    let other = Arc::new(Worker::new(spec("b", Stage::Plain, "m")));
    let _busy = other.load_guard();
    for engines in [vec![engine.clone()], vec![other.clone(), engine.clone()]] {
        for (reject, invalid) in [(false, false), (true, false), (false, true)] {
            let check = Arc::new(Check {
                calls: Mutex::new(Vec::new()),
                reject,
                invalid,
            });
            let mut policy = PowerOfTwoPolicy::new(EngineReportedLoadTable::new());
            policy.admission = check.clone();
            assert!(matches!(
                policy.pick(&[], &request).await,
                Err(PickError::NoCandidates)
            ));
            assert!(check.calls.lock().unwrap().is_empty());
            let result = policy.pick(&engines, &request).await;
            if invalid {
                assert!(matches!(result, Err(PickError::InvalidSignal(_))));
            } else if reject {
                assert!(matches!(result, Err(PickError::AdmissionRejected(reason))
                if reason.engine == engine.id && reason.reason == "full"));
            } else {
                assert!(Arc::ptr_eq(&result.unwrap().engine, &engine));
            }
            assert_eq!(
                check.calls.lock().unwrap().as_slice(),
                std::slice::from_ref(&engine.id)
            );
        }
    }
 }
@@ -0,0 +1,138 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
 use sgl_router::policies_reorg::admission::{Decision, EngineAdmission};
 use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
 use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
 use sgl_router::state::load_monitor::engine_reported_load::{
    EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat,
 };
 use sgl_router::workers::Worker;
 const URL: &str = "http://engine";
 fn report(table: &EngineReportedLoadTable, rank: u32, running: u64, at: Instant) {
    table.set(
        URL,
        rank,
        LoadStat {
            num_running_reqs: running,
            num_waiting_reqs: 2,
            num_tokens: 30,
            max_total_num_tokens: 100,
            native_cache: None,
        },
        at,
    );
 }
 fn engine() -> Arc<Worker> {
    Arc::new(Worker::new(WorkerSpec {
        id: WorkerId("a".into()),
        url: URL.into(),
        mode: Stage::Plain,
        model_ids: vec![ModelId("m".into())],
        bootstrap_port: None,
    }))
 }
 #[derive(Debug)]
 struct ObserveAdmission {
    table: Arc<EngineReportedLoadTable>,
    observations: Mutex<Vec<Option<EngineReportedWorkerLoad>>>,
 }
 impl EngineAdmission for ObserveAdmission {
    fn check(
        &self,
        engine: &Worker,
        _: &PickRequest<'_>,
        load: Option<&EngineReportedWorkerLoad>,
    ) -> Result<Decision, PickError> {
        assert_eq!(engine.url, URL);
        // A new report arriving after selection must not change the observation
        // supplied to admission. The next pick should read the new report.
        report(&self.table, 0, 99, Instant::now());
        self.observations.lock().unwrap().push(load.cloned());
        Ok(Decision::Allow)
    }
 }
 #[tokio::test]
 async fn selected_load_reaches_admission_and_next_pick_reads_fresh_state() {
    let table = EngineReportedLoadTable::new();
    let first_at = Instant::now();
    report(&table, 0, 1, first_at);
    report(&table, 1, 3, first_at);
    table.mark_expected_rank(URL, 0);
    table.mark_expected_rank(URL, 1);
    let admission = Arc::new(ObserveAdmission {
        table: table.clone(),
        observations: Mutex::default(),
    });
    let mut policy = PowerOfTwoPolicy::new(table);
    policy.admission = admission.clone();
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    let alternative = Arc::new(Worker::new(WorkerSpec {
        id: WorkerId("b".into()),
        url: "http://other".into(),
        mode: Stage::Plain,
        model_ids: vec![ModelId("m".into())],
        bootstrap_port: None,
    }));
    // These old-format reports lack native pressure metrics, so selection uses
    // local active counts for both candidates and chooses the second engine.
    let _busy = alternative.load_guard();
    let engines = [alternative, engine()];
    for _ in 0..2 {
        let pick = policy.pick(&engines, &request).await.unwrap();
        assert!(Arc::ptr_eq(&pick.engine, &engines[1]));
    }
    let observations = admission.observations.lock().unwrap();
    assert_eq!(observations.len(), 2);
    assert_eq!(
        observations[0],
        Some(EngineReportedWorkerLoad {
            num_running_reqs: 4,
            num_waiting_reqs: 4,
            num_tokens: 60,
            max_total_num_tokens: 200,
            captured_at: first_at,
        })
    );
    assert_eq!(observations[1].as_ref().unwrap().num_running_reqs, 102);
 }
 #[tokio::test]
 async fn missing_stale_and_incomplete_reports_reach_admission_as_unknown() {
    for case in ["missing", "stale", "incomplete"] {
        let table = EngineReportedLoadTable::new();
        match case {
            "missing" => {}
            "stale" => report(&table, 0, 1, Instant::now() - Duration::from_secs(3600)),
            "incomplete" => {
                report(&table, 0, 1, Instant::now());
                table.mark_expected_rank(URL, 0);
                table.mark_expected_rank(URL, 1);
            }
            _ => unreachable!(),
        }
        let admission = Arc::new(ObserveAdmission {
            table: table.clone(),
            observations: Mutex::default(),
        });
        let mut policy = PowerOfTwoPolicy::new(table);
        policy.admission = admission.clone();
        let model = ModelId("m".into());
        policy
            .pick(&[engine()], &PickRequest::new(&model, Stage::Plain, 10))
            .await
            .unwrap();
        assert_eq!(*admission.observations.lock().unwrap(), [None], "{case}");
    }
 }
@@ -0,0 +1,203 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
 use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
 use sgl_router::policies_reorg::{PickRequest, Policy, Stage};
 use sgl_router::state::load_monitor::engine_reported_load::{
    EngineReportedLoadTable, LoadStat, NativeCacheRankLoad,
 };
 use sgl_router::workers::Worker;
 fn engine(id: &str, stage: Stage, active: usize) -> Arc<Worker> {
    let worker = Arc::new(Worker::new(WorkerSpec {
        id: WorkerId(id.into()),
        url: format!("http://{id}"),
        mode: stage,
        model_ids: vec![ModelId("m".into())],
        bootstrap_port: None,
    }));
    worker.active_requests.store(active, Ordering::Relaxed);
    worker
 }
 fn load(running: u64, waiting: u64, tokens: u64, capacity: u64, pending: u64) -> LoadStat {
    LoadStat {
        num_running_reqs: running,
        num_waiting_reqs: waiting,
        num_tokens: tokens,
        max_total_num_tokens: capacity,
        native_cache: Some(NativeCacheRankLoad {
            num_waiting_uncached_tokens: pending,
            num_total_tokens: tokens,
            max_running_requests: 100,
            total_prefill_uncached_tokens: 0,
            total_prefill_busy_us: 0,
        }),
    }
 }
 async fn assert_winner(
    policy: &PowerOfTwoPolicy,
    engines: &[Arc<Worker>],
    stage: Stage,
    expected: &Arc<Worker>,
 ) {
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, stage, 10);
    // With exactly two candidates the winner is independent of sample order.
    for _ in 0..16 {
        let pick = policy.pick(engines, &request).await.unwrap();
        assert!(Arc::ptr_eq(&pick.engine, expected), "stage: {stage:?}");
        assert_eq!(pick.reason, "power_of_two");
    }
 }
 #[tokio::test]
 async fn plain_and_prefill_use_pending_work_while_decode_uses_request_pressure() {
    for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] {
        let engines = [engine("a", stage, 100), engine("b", stage, 0)];
        let table = EngineReportedLoadTable::new();
        table.set(&engines[0].url, 0, load(5, 8, 10, 100, 1), Instant::now());
        table.set(&engines[1].url, 0, load(1, 1, 10, 100, 100), Instant::now());
        let expected = if stage == Stage::Decode { 1 } else { 0 };
        assert_winner(
            &PowerOfTwoPolicy::new(table),
            &engines,
            stage,
            &engines[expected],
        )
        .await;
    }
 }
 #[tokio::test]
 async fn prefill_uses_estimated_queue_time_only_when_both_engines_have_rates() {
    for both_have_rates in [true, false] {
        let engines = [
            engine("a", Stage::Prefill, 0),
            engine("b", Stage::Prefill, 0),
        ];
        let table = EngineReportedLoadTable::new();
        for (i, worker) in engines.iter().enumerate() {
            let mut report = load(1, 1, 10, 100, if i == 0 { 10 } else { 20 });
            if i == 0 || both_have_rates {
                table.set(&worker.url, 0, report.clone(), Instant::now());
            }
            let native = report.native_cache.as_mut().unwrap();
            native.total_prefill_uncached_tokens = if i == 0 { 100 } else { 1000 };
            native.total_prefill_busy_us = 1_000_000;
            table.set(&worker.url, 0, report, Instant::now());
        }
        // B has more queued tokens, but its higher throughput gives a shorter
        // estimated queue. Without B's rate, compare queued tokens for both.
        let expected = usize::from(both_have_rates);
        assert_winner(
            &PowerOfTwoPolicy::new(table),
            &engines,
            Stage::Prefill,
            &engines[expected],
        )
        .await;
    }
 }
 #[tokio::test]
 async fn decode_orders_by_waiting_running_kv_fraction_then_tokens() {
    let cases = [
        (load(50, 1, 90, 100, 0), load(1, 2, 1, 100, 0)),
        (load(1, 1, 90, 100, 0), load(2, 1, 1, 100, 0)),
        (load(1, 1, 100, 1000, 0), load(1, 1, 20, 100, 0)),
        (load(1, 1, 10, 100, 0), load(1, 1, 100, 1000, 0)),
    ];
    for (left, right) in cases {
        let engines = [
            engine("a", Stage::Decode, 100),
            engine("b", Stage::Decode, 0),
        ];
        let table = EngineReportedLoadTable::new();
        table.set(&engines[0].url, 0, left, Instant::now());
        table.set(&engines[1].url, 0, right, Instant::now());
        assert_winner(
            &PowerOfTwoPolicy::new(table),
            &engines,
            Stage::Decode,
            &engines[0],
        )
        .await;
    }
 }
 #[tokio::test]
 async fn unusable_telemetry_falls_back_to_local_load_for_both_candidates() {
    for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] {
        for case in [
            "missing",
            "stale",
            "incomplete",
            "old_publisher",
            "unknown_capacity",
        ] {
            let engines = [engine("a", stage, 1), engine("b", stage, 5)];
            let table = EngineReportedLoadTable::new();
            // A's high reported pressure must not be compared to B's local
            // count or to a fabricated zero for its unavailable telemetry.
            table.set(
                &engines[0].url,
                0,
                load(90, 90, 90, 100, 900),
                Instant::now(),
            );
            let mut right = load(0, 0, 0, 100, 0);
            let mut at = Instant::now();
            match case {
                "missing" => {}
                "stale" => at -= Duration::from_secs(3600),
                "incomplete" => {
                    table.mark_expected_rank(&engines[1].url, 0);
                    table.mark_expected_rank(&engines[1].url, 1);
                }
                "old_publisher" => right.native_cache = None,
                "unknown_capacity" => right.max_total_num_tokens = 0,
                _ => unreachable!(),
            }
            if case != "missing" {
                table.set(&engines[1].url, 0, right, at);
            }
            assert_winner(&PowerOfTwoPolicy::new(table), &engines, stage, &engines[0]).await;
        }
    }
 }
 #[tokio::test]
 async fn equal_reported_pressure_uses_local_active_load_as_tiebreaker() {
    for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] {
        let engines = [engine("a", stage, 5), engine("b", stage, 1)];
        let table = EngineReportedLoadTable::new();
        for worker in &engines {
            table.set(&worker.url, 0, load(1, 1, 10, 100, 10), Instant::now());
        }
        assert_winner(&PowerOfTwoPolicy::new(table), &engines, stage, &engines[1]).await;
    }
 }
 #[tokio::test]
 async fn multiple_candidates_never_select_the_unique_busiest_engine() {
    let engines: Vec<_> = (0..8)
        .map(|i| engine(&i.to_string(), Stage::Plain, i))
        .collect();
    let policy = PowerOfTwoPolicy::new(EngineReportedLoadTable::new());
    let model = ModelId("m".into());
    let request = PickRequest::new(&model, Stage::Plain, 10);
    for _ in 0..64 {
        let pick = policy.pick(&engines, &request).await.unwrap();
        // Every distinct pair has an engine less busy than the last candidate.
        assert!(engines[..7]
            .iter()
            .any(|engine| Arc::ptr_eq(engine, &pick.engine)));
    }
 }
@@ -21,6 +21,8 @@ use std::sync::Arc;
 use std::time::Duration;
 use tower::ServiceExt;
 mod reorg;
 const TEST_TIMEOUT: Duration = Duration::from_secs(5);
 fn config_for(_worker_url: &str) -> Config {
@@ -0,0 +1,523 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 use super::*;
 use crate::common::mock_worker::MockWorker;
 use futures::future::BoxFuture;
 use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup};
 use sgl_router::policies::PolicyRegistry;
 use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission};
 use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
 use sgl_router::server::app_context::ChatRouting;
 use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
 use std::sync::Mutex;
 type PickCall = (String, Stage, u64, Option<u64>);
 #[derive(Debug)]
 struct FirstPolicy {
    calls: Mutex<Vec<PickCall>>,
    admission: Arc<dyn EngineAdmission>,
    miss: bool,
    invalid: bool,
 }
 impl Default for FirstPolicy {
    fn default() -> Self {
        Self {
            admission: Arc::new(AllowAll),
            miss: false,
            invalid: false,
            calls: Mutex::new(Vec::new()),
        }
    }
 }
 impl Policy for FirstPolicy {
    fn pick<'a>(
        &'a self,
        engines: &'a [Arc<Worker>],
        request: &'a PickRequest<'a>,
    ) -> BoxFuture<'a, Result<Pick, PickError>> {
        Box::pin(async move {
            self.calls.lock().unwrap().push((
                request.bucket.to_owned(),
                request.stage,
                request.input_tokens,
                request.expected_peak_tokens,
            ));
            if self.invalid {
                return Err(PickError::InvalidSignal("invalid policy input".into()));
            }
            if self.miss {
                return Err(PickError::NoCandidates);
            }
            if engines.is_empty() {
                return Err(PickError::NoCandidates);
            }
            let engine = engines[0].clone();
            if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? {
                return Err(PickError::AdmissionRejected(Rejection {
                    engine: engine.id.clone(),
                    reason,
                }));
            }
            Ok(Pick {
                engine,
                reason: "test",
            })
        })
    }
 }
 #[derive(Debug)]
 struct RejectAll;
 impl EngineAdmission for RejectAll {
    fn check(
        &self,
        _: &Worker,
        _: &PickRequest<'_>,
        _: Option<&EngineReportedWorkerLoad>,
    ) -> Result<Decision, PickError> {
        Ok(Decision::Reject("full".into()))
    }
 }
 fn rejecting_policy() -> Arc<FirstPolicy> {
    Arc::new(FirstPolicy {
        admission: Arc::new(RejectAll),
        ..Default::default()
    })
 }
 fn group(id: &str, policy: Arc<FirstPolicy>) -> EngineGroup {
    EngineGroup {
        worker_ids: Some([WorkerId(id.into())].into_iter().collect()),
        policy,
    }
 }
 fn context(workers: &[(&str, Stage, &MockWorker)], buckets: Vec<Bucket>) -> Arc<AppContext> {
    let config = config_for("");
    let registry = Arc::new(WorkerRegistry::default());
    for &(id, mode, worker) in workers {
        registry
            .add(WorkerSpec {
                id: WorkerId(id.into()),
                url: worker.url.clone(),
                mode,
                model_ids: vec![ModelId("tiny".into())],
                bootstrap_port: Some(8998),
            })
            .unwrap();
    }
    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap());
    let mut ctx = AppContext::new(
        config,
        tokenizers,
        Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()),
        registry,
        // The reorg route must not require the legacy policy registry.
        Arc::new(PolicyRegistry::default()),
    );
    ctx.chat_routing = ChatRouting::Reorg(
        [(ModelId("tiny".into()), BucketResolver::new(buckets))]
            .into_iter()
            .collect(),
    );
    Arc::new(ctx)
 }
 fn request(value: serde_json::Value) -> Request<Body> {
    Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_vec(&value).unwrap()))
        .unwrap()
 }
 fn body(content: &str) -> serde_json::Value {
    serde_json::json!({"model": "tiny", "messages": [{"role": "user", "content": content}]})
 }
 #[tokio::test]
 async fn length_selects_plain_bucket_before_engine_selection() {
    let short_worker = MockWorker::start(vec![]).await;
    let long_worker = MockWorker::start(vec![]).await;
    let policy = Arc::new(FirstPolicy::default());
    let mut short = Bucket::new("short", BucketGroups::Plain(group("short", policy.clone())));
    short.limits.max = Some(4);
    let long = Bucket::new("long", BucketGroups::Plain(group("long", policy.clone())));
    let ctx = context(
        &[
            ("short", Stage::Plain, &short_worker),
            ("long", Stage::Plain, &long_worker),
        ],
        vec![long, short],
    );
    let app = build_router(ctx);
    let response = app.clone().oneshot(request(body("hi"))).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let _ = response.into_body().collect().await.unwrap();
    assert!(short_worker.captured.lock().unwrap().last_body.is_some());
    assert!(long_worker.captured.lock().unwrap().last_body.is_none());
    let response = app
        .oneshot(request(body(&"hello ".repeat(30))))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let _ = response.into_body().collect().await.unwrap();
    assert!(long_worker.captured.lock().unwrap().last_body.is_some());
    let calls = policy.calls.lock().unwrap();
    assert_eq!(calls.len(), 2);
    assert_eq!((&*calls[0].0, calls[0].1), ("short", Stage::Plain));
    assert_eq!((&*calls[1].0, calls[1].1), ("long", Stage::Plain));
 }
 #[tokio::test]
 async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() {
    let prefill = MockWorker::start(vec![]).await;
    let decode = MockWorker::start(vec![]).await;
    let policy = Arc::new(FirstPolicy::default());
    let mut selected = Bucket::new(
        "selected",
        BucketGroups::Pd {
            prefill: group("p", policy.clone()),
            decode: group("d", policy.clone()),
        },
    );
    selected.limits.max = Some(100);
    let other_policy = Arc::new(FirstPolicy::default());
    let other = Bucket::new(
        "other",
        BucketGroups::Pd {
            prefill: group("p", other_policy.clone()),
            decode: group("d", other_policy.clone()),
        },
    );
    let ctx = context(
        &[
            ("p", Stage::Prefill, &prefill),
            ("d", Stage::Decode, &decode),
        ],
        vec![other, selected],
    );
    let mut body = body("hello");
    body["max_completion_tokens"] = 10.into();
    let response = build_router(ctx).oneshot(request(body)).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(response.headers()["x-sgl-decode-url"], decode.url);
    let _ = response.into_body().collect().await.unwrap();
    tokio::time::timeout(TEST_TIMEOUT, async {
        while prefill.captured.lock().unwrap().last_body.is_none() {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    let p: serde_json::Value =
        serde_json::from_slice(prefill.captured.lock().unwrap().last_body.as_ref().unwrap())
            .unwrap();
    let d: serde_json::Value =
        serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap())
            .unwrap();
    assert!(p["bootstrap_room"].is_number());
    assert_eq!(p["bootstrap_room"], d["bootstrap_room"]);
    let calls = policy.calls.lock().unwrap();
    assert_eq!(calls.len(), 2);
    assert_eq!((&*calls[0].0, calls[0].1), ("selected", Stage::Prefill));
    assert_eq!((&*calls[1].0, calls[1].1), ("selected", Stage::Decode));
    assert_eq!(calls[0].3, Some(calls[0].2 + 10));
    assert_eq!(calls[1].3, calls[0].3);
    assert!(other_policy.calls.lock().unwrap().is_empty());
 }
 #[tokio::test]
 async fn missing_decode_in_all_buckets_does_not_dispatch_prefill() {
    let prefill = MockWorker::start(vec![]).await;
    let decode = MockWorker::start(vec![]).await;
    let policy = Arc::new(FirstPolicy::default());
    let mut selected = Bucket::new(
        "selected",
        BucketGroups::Pd {
            prefill: group("p", policy.clone()),
            decode: group("missing", policy.clone()),
        },
    );
    selected.limits.max = Some(100);
    let other = Bucket::new(
        "other",
        BucketGroups::Pd {
            prefill: group("p", policy.clone()),
            decode: group("also-missing", policy.clone()),
        },
    );
    let ctx = context(
        &[
            ("p", Stage::Prefill, &prefill),
            ("d", Stage::Decode, &decode),
        ],
        vec![other, selected],
    );
    let response = build_router(ctx.clone())
        .oneshot(request(body("hi")))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
    assert_eq!(
        response.headers()["x-router-error-code"],
        "no_decode_workers_available"
    );
    assert!(prefill.captured.lock().unwrap().last_body.is_none());
    assert!(decode.captured.lock().unwrap().last_body.is_none());
    assert_eq!(policy.calls.lock().unwrap().len(), 2);
    assert_eq!(ctx.router_inflight_load.inflight_count(), 0);
    assert_eq!(
        ctx.registry
            .get(&WorkerId("p".into()))
            .unwrap()
            .router_inflight_load(),
        0
    );
 }
 #[tokio::test]
 async fn rejects_unsupported_length_unknown_model_and_overflow_before_policy() {
    let worker = MockWorker::start(vec![]).await;
    let policy = Arc::new(FirstPolicy::default());
    let mut bucket = Bucket::new("short", BucketGroups::Plain(group("w", policy.clone())));
    bucket.max_context_tokens = Some(4);
    let app = build_router(context(&[("w", Stage::Plain, &worker)], vec![bucket]));
    let mut long = body("hi");
    long["max_tokens"] = 100.into();
    let mut overflow = body("hi");
    overflow["max_tokens"] = u64::MAX.into();
    let mut unknown = body("hi");
    unknown["model"] = "unknown".into();
    for (body, status) in [
        (long, StatusCode::BAD_REQUEST),
        (overflow, StatusCode::BAD_REQUEST),
        (unknown, StatusCode::NOT_FOUND),
        (serde_json::json!({}), StatusCode::BAD_REQUEST),
    ] {
        let response = app.clone().oneshot(request(body)).await.unwrap();
        assert_eq!(response.status(), status);
    }
    assert!(policy.calls.lock().unwrap().is_empty());
    assert!(worker.captured.lock().unwrap().last_body.is_none());
 }
 #[tokio::test]
 async fn streaming_uses_existing_forwarder() {
    let worker = MockWorker::start(vec!["data: {\"choices\":[]}\n\n", "data: [DONE]\n\n"]).await;
    let policy = Arc::new(FirstPolicy::default());
    let bucket = Bucket::new("plain", BucketGroups::Plain(group("w", policy)));
    let app = build_router(context(&[("w", Stage::Plain, &worker)], vec![bucket]));
    let mut body = body("hi");
    body["stream"] = true.into();
    let response = app.oneshot(request(body)).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    assert!(response.headers()["content-type"]
        .to_str()
        .unwrap()
        .starts_with("text/event-stream"));
    let bytes = response.into_body().collect().await.unwrap().to_bytes();
    assert!(String::from_utf8_lossy(&bytes).contains("data: [DONE]"));
 }
 #[tokio::test]
 async fn reorg_route_keeps_chat_body_limit() {
    let app = build_router(context(&[], vec![]));
    let request = Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .body(Body::from(vec![b' '; MAX_CHAT_BODY_BYTES + 1]))
        .unwrap();
    assert_eq!(
        app.oneshot(request).await.unwrap().status(),
        StatusCode::PAYLOAD_TOO_LARGE
    );
 }
 #[tokio::test]
 async fn plain_fallback_skips_empty_missed_and_rejected_buckets_then_stops_on_success() {
    let rejected_worker = MockWorker::start(vec![]).await;
    let winner = MockWorker::start(vec![]).await;
    let skipped = Arc::new(FirstPolicy::default());
    let rejected = rejecting_policy();
    let missed = Arc::new(FirstPolicy {
        miss: true,
        ..Default::default()
    });
    let accepted = Arc::new(FirstPolicy::default());
    let buckets = vec![
        Bucket::new(
            "a-empty",
            BucketGroups::Plain(group("missing", skipped.clone())),
        ),
        Bucket::new(
            "b-miss",
            BucketGroups::Plain(group("rejected", missed.clone())),
        ),
        Bucket::new(
            "c-rejected",
            BucketGroups::Plain(group("rejected", rejected.clone())),
        ),
        Bucket::new(
            "d-winner",
            BucketGroups::Plain(group("winner", accepted.clone())),
        ),
        Bucket::new(
            "e-unused",
            BucketGroups::Plain(group("winner", skipped.clone())),
        ),
    ];
    let ctx = context(
        &[
            ("rejected", Stage::Plain, &rejected_worker),
            ("winner", Stage::Plain, &winner),
        ],
        buckets,
    );
    let response = build_router(ctx)
        .oneshot(request(body("hi")))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let _ = response.into_body().collect().await.unwrap();
    assert!(rejected_worker.captured.lock().unwrap().last_body.is_none());
    assert!(winner.captured.lock().unwrap().last_body.is_some());
    assert!(skipped.calls.lock().unwrap().is_empty());
    assert_eq!(missed.calls.lock().unwrap().len(), 1);
    assert_eq!(rejected.calls.lock().unwrap().len(), 1);
    assert_eq!(accepted.calls.lock().unwrap()[0].0, "d-winner");
 }
 #[tokio::test]
 async fn decode_failure_retries_both_groups_in_next_bucket_without_dispatching_first_prefill() {
    // Cover an empty decode group and rejection of a selected decode engine.
    for reject_decode in [false, true] {
        let first_prefill = MockWorker::start(vec![]).await;
        let second_prefill = MockWorker::start(vec![]).await;
        let decode = MockWorker::start(vec![]).await;
        let first = Arc::new(FirstPolicy::default());
        let rejected = rejecting_policy();
        let accepted = Arc::new(FirstPolicy::default());
        let buckets = vec![
            Bucket::new(
                "a-first",
                BucketGroups::Pd {
                    prefill: group("p1", first.clone()),
                    decode: group(if reject_decode { "d" } else { "missing" }, rejected),
                },
            ),
            Bucket::new(
                "b-second",
                BucketGroups::Pd {
                    prefill: group("p2", accepted.clone()),
                    decode: group("d", accepted.clone()),
                },
            ),
        ];
        let ctx = context(
            &[
                ("p1", Stage::Prefill, &first_prefill),
                ("p2", Stage::Prefill, &second_prefill),
                ("d", Stage::Decode, &decode),
            ],
            buckets,
        );
        let response = build_router(ctx.clone())
            .oneshot(request(body("hi")))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let _ = response.into_body().collect().await.unwrap();
        tokio::time::timeout(TEST_TIMEOUT, async {
            while second_prefill.captured.lock().unwrap().last_body.is_none() {
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();
        assert!(first_prefill.captured.lock().unwrap().last_body.is_none());
        assert!(decode.captured.lock().unwrap().last_body.is_some());
        assert_eq!(
            ctx.registry
                .get(&WorkerId("p1".into()))
                .unwrap()
                .router_inflight_load(),
            0
        );
        assert_eq!(first.calls.lock().unwrap().len(), 1);
        let calls = accepted.calls.lock().unwrap();
        assert_eq!(calls.len(), 2);
        assert_eq!((&*calls[0].0, calls[0].1), ("b-second", Stage::Prefill));
        assert_eq!((&*calls[1].0, calls[1].1), ("b-second", Stage::Decode));
    }
 }
 #[tokio::test]
 async fn admission_exhaustion_is_preserved_when_later_buckets_are_empty() {
    let worker = MockWorker::start(vec![]).await;
    let first = rejecting_policy();
    let second = rejecting_policy();
    let empty = Arc::new(FirstPolicy::default());
    let ctx = context(
        &[("w", Stage::Plain, &worker)],
        vec![
            Bucket::new("a-first", BucketGroups::Plain(group("w", first.clone()))),
            Bucket::new("b-second", BucketGroups::Plain(group("w", second.clone()))),
            Bucket::new(
                "c-empty",
                BucketGroups::Plain(group("missing", empty.clone())),
            ),
        ],
    );
    let response = build_router(ctx.clone())
        .oneshot(request(body("hi")))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
    assert_eq!(
        response.headers()["x-router-error-code"],
        "policy_selection_failed"
    );
    assert_eq!(first.calls.lock().unwrap().len(), 1);
    assert_eq!(second.calls.lock().unwrap().len(), 1);
    assert!(empty.calls.lock().unwrap().is_empty());
    assert!(worker.captured.lock().unwrap().last_body.is_none());
    assert_eq!(ctx.router_inflight_load.inflight_count(), 0);
 }
 #[tokio::test]
 async fn invalid_policy_signal_stops_bucket_iteration() {
    let worker = MockWorker::start(vec![]).await;
    let invalid = Arc::new(FirstPolicy {
        invalid: true,
        ..Default::default()
    });
    let later = Arc::new(FirstPolicy::default());
    let ctx = context(
        &[("w", Stage::Plain, &worker)],
        vec![
            Bucket::new(
                "a-invalid",
                BucketGroups::Plain(group("w", invalid.clone())),
            ),
            Bucket::new("b-later", BucketGroups::Plain(group("w", later.clone()))),
        ],
    );
    let response = build_router(ctx)
        .oneshot(request(body("hi")))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
    assert_eq!(invalid.calls.lock().unwrap().len(), 1);
    assert!(later.calls.lock().unwrap().is_empty());
    assert!(worker.captured.lock().unwrap().last_body.is_none());
 }