feat(router): implement ThunderAgent request classifier

Cargo.toml

@@ -12,3 +12,6 @@ version = "0.1.0"
 edition = "2024"
 license = "Apache-2.0"
 repository = "https://github.com/ishandhanani/router-sandbox"
+
+[workspace.dependencies]
+dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", rev = "05747ac66aa745160a8c223683efc4597b4868fb" }

README.md

@@ -4,7 +4,7 @@ This Cargo workspace contains external routing strategies for Dynamo.
 
 ## Policy crates
 
-- [`thunderagent-dynamo-policy`](crates/thunderagent/README.md) provides session-aware admission and worker selection.
+- [`thunderagent-dynamo-policy`](crates/thunderagent/README.md) provides program-aware request classification and flow control.
 - [`sticky-until-saturated-dynamo-policy`](crates/sticky-until-saturated/README.md) provides token-aware cache-affinity routing.
 - [`sglang-cache-aware-dynamo-policy`](crates/sglang-cache-aware/README.md) ports SGLang's cache-aware load gate using Dynamo's authoritative device-KV overlap.
 - [`sgl-router-cache-aware-dynamo-policy`](crates/sgl-router-cache-aware/README.md) ports the experimental SGL Router cache-aware policy used in the router comparison.

crates/sgl-router-cache-aware/Cargo.toml

@@ -10,4 +10,4 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
-dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", branch = "main", features = ["standalone-selection"] }
+dynamo-kv-router = { workspace = true, features = ["standalone-selection"] }

crates/sglang-cache-aware/Cargo.toml

@@ -10,6 +10,6 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
-dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", branch = "main", features = ["standalone-selection"] }
+dynamo-kv-router = { workspace = true, features = ["standalone-selection"] }
 serde = { version = "1", features = ["derive"] }
 thiserror = "2"

crates/sticky-until-saturated/Cargo.toml

@@ -13,7 +13,7 @@ repository.workspace = true
 bench = []
 
 [dependencies]
-dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", branch = "main", features = ["standalone-selection"] }
+dynamo-kv-router = { workspace = true, features = ["standalone-selection"] }
 serde = { version = "1", features = ["derive"] }
 thiserror = "2"
 

crates/thunderagent/Cargo.toml

@@ -1,24 +1,20 @@
 [package]
 name = "thunderagent-dynamo-policy"
-description = "ThunderAgent session-aware admission and routing policy for Dynamo"
+description = "ThunderAgent program-aware request classifier for Dynamo"
 version.workspace = true
 edition.workspace = true
 license.workspace = true
 repository.workspace = true
 
 [dependencies]
-dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", branch = "idhanani/dyn-3231-session-aware-policy-crate", features = ["standalone-selection"] }
+async-trait = "0.1"
+dynamo-kv-router.workspace = true
 parking_lot = "0.12"
 serde = { version = "1", features = ["derive"] }
 thiserror = "2"
+tokio = { version = "1", features = ["sync", "time"] }
 tracing = "0.1"
 
 [dev-dependencies]
-criterion = "0.5"
-dynamo-kv-router = { git = "https://github.com/ai-dynamo/dynamo.git", branch = "idhanani/dyn-3231-session-aware-policy-crate", features = ["bench", "standalone-selection"] }
-tokio = { version = "1", features = ["rt-multi-thread"] }
-tokio-util = "0.7"
-
-[[bench]]
-name = "admission_overhead"
-harness = false
+tempfile = "3"
+tokio = { version = "1", features = ["macros", "rt", "sync", "time"] }

crates/thunderagent/README.md

@@ -1,69 +1,110 @@
 # ThunderAgent
 
-`thunderagent-dynamo-policy` is a session-aware Dynamo routing policy. It owns the ThunderAgent admission and capacity state outside Dynamo while Dynamo continues to own request storage and serially invokes the policy from its queue actor.
+`thunderagent-dynamo-policy` implements ThunderAgent's program-aware flow control on Dynamo's asynchronous request-classifier API. It registers a request classifier and a worker-selection policy through Dynamo's statically linked router-plugin catalog; no Python adapter is required.
+
+This prototype uses the classifier signals in [ai-dynamo/dynamo#14123](https://github.com/ai-dynamo/dynamo/pull/14123) and the unified plugin catalog in [#14149](https://github.com/ai-dynamo/dynamo/pull/14149). One `register(&mut RouterPluginRegistry)` call registers both roles through `plugins::request_classifier` and `plugins::worker_selection`.
+
+## Source layout
+
+```text
+src/
+├── lib.rs                       # Public exports and catalog registration façade
+├── config.rs                    # ThunderAgent configuration and validation
+├── request_classifier/
+│   ├── mod.rs                   # RequestClassifier implementation and plugin factory
+│   ├── capacity.rs              # Cached worker capacity and liveness view
+│   └── scheduler.rs             # Program state, admission, pause, resume, and repacking
+└── worker_selection/
+    └── mod.rs                   # Stateless load scorer, target-aware picker, and plugin factory
+```
 
-Runnable ordering remains Dynamo-owned. This prototype depends on the Dynamo queue-admission seam in [ai-dynamo/dynamo#13019](https://github.com/ai-dynamo/dynamo/pull/13019).
+## Ownership and routing
 
-## Algorithm
+`ThunderAgentClassifier` owns the complete program table: per-session serialization, reasoning/acting state, pause state, retained token usage, and worker assignment. The state is internally reference-counted because `classify` must return independently pollable `'static` futures, but it is not shared with the worker selector.
 
-The returned `WorkerSelectionPolicy` combines three components:
+The classifier chooses a request-local worker target when it releases a request. ThunderAgent's stateless selector honors that target when it is eligible and otherwise falls back to current least-loaded placement. The `Sent` lifecycle event reconciles the program table with the worker Dynamo actually selected. If a worker disappears, the classifier clears the assignment and chooses a live replacement.
 
-1. The queue-admission policy serializes each session: a session has at most one runnable request, and subsequent requests wait in session order.
-2. For a session without an existing assignment, admission selects an eligible available worker with enough reported token capacity. It chooses the least-used worker that can fit the request context plus `buffer_per_program`.
-3. The picker preserves that session's assigned worker when it remains eligible; otherwise it falls back to the minimum Dynamo load score, breaking ties by stable worker key.
+Use soft session affinity. Affinity preserves locality between requests, while soft mode allows the classifier to repack a paused program onto a different worker. Hard affinity would prevent that Python-parity behavior before the custom selector runs.
 
-Requests without Dynamo session context bypass the session admission logic and use the fallback scorer and picker directly.
+Requests without Dynamo session context return from classification immediately.
 
 ```mermaid
 flowchart TD
-    A["Request arrives"] --> B{"Has session context?"}
-    B -- "No" --> C["Route by fallback load score"]
-    B -- "Yes" --> D{"Another request is active for the session?"}
-    D -- "Yes" --> E["Defer in session FIFO queue"]
-    D -- "No" --> F{"Existing eligible assignment?"}
-    F -- "Yes" --> G["Mark request runnable"]
-    F -- "No" --> H["Select least-used worker with enough capacity"]
-    H --> I{"Capacity available?"}
-    I -- "Yes" --> G
-    I -- "No" --> J["Pause and defer session"]
-    E --> K["Periodic capacity reconciliation"]
-    J --> K
-    K --> L["Resume eligible sessions or force after timeout"]
+    A["Request enters classify"] --> B{"Has session context?"}
+    B -- "No" --> C["Return unchanged"]
+    B -- "Yes" --> D["Register in classifier-owned program table"]
+    D --> E{"Session busy or program paused?"}
+    E -- "Yes" --> F["Keep the classify future pending"]
+    E -- "No" --> G{"Known pressure permits release?"}
+    G -- "No" --> F
+    G -- "Yes" --> H["Attach chosen worker target and return request"]
+    H --> I["Stateless selector honors target or falls back"]
+    I --> J["Sent reconciles the actual worker"]
+    K["Completed or Aborted"] --> L["Commit or roll back program state"]
+    L --> M["Notify pending classifier futures"]
+    M --> F
 ```
 
-On each reconciliation interval, the policy accounts for active program tokens plus the configured per-program buffer. Acting programs use a configurable token weight and a decayed estimate for forced timeout resumption. When a worker exceeds `pause_threshold`, the policy pauses smaller acting programs first until usage reaches `pause_target`; reasoning programs are marked to pause after their current request completes. It resumes paused sessions greedily below `pause_threshold - resume_hysteresis`, preserving an existing assignment when possible. A paused session is force-resumed after `resume_timeout_seconds` when an eligible worker is available.
+## Score and pick
+
+Worker selection is a stateless execution step for the classifier's placement decision. It does not read or mutate the classifier's program table, decide when to pause or resume, or own session pinning.
+
+The scorer requires Dynamo's worker `LOAD` input and assigns each eligible candidate this lower-is-better cost:
+
+```text
+active_prefill_tokens + (decode_cost_blocks * kv_cache_block_size)
+```
+
+The picker then applies two rules in order:
+
+1. If the request carries a classifier or session-affinity target and that worker/rank remains eligible, select it regardless of its load score. This is how the classifier's best-fit-decreasing placement is executed.
+2. Otherwise, select the eligible candidate with the lowest score, breaking ties by worker identity for deterministic behavior. This handles a removed, unhealthy, or otherwise ineligible target without failing the request.
 
-Completed sessions transition to an acting state and retain their assignment for `session_retention_seconds`; inactive retained sessions are later expired. Worker removal clears affected assignments.
+The target is advisory, not a reservation. Dynamo applies routing constraints and worker eligibility before the picker runs. After dispatch, the `Sent` classifier event reports the worker/rank actually selected so the classifier can reconcile its assignment and capacity accounting. Soft session affinity preserves normal cross-turn locality but permits the classifier to clear or replace the request-local target when it pauses and repacks a program.
+
+The catalog factory receives a cached view of Dynamo's existing discovery state. It derives each worker/rank's program-retention budget as `kv_cache_block_size * total_kv_blocks`, while representing worker liveness separately from missing capacity metadata. The callback only reads the cached view; it performs no discovery, model-card parsing, or blocking I/O on the classification path.
+
+ThunderAgent computes live used capacity from its own program table. On each reconciliation, it accounts for active program tokens plus `buffer_per_program`. Acting programs use `acting_token_weight`. When a worker exceeds `pause_threshold`, the classifier pauses smaller acting programs first until usage reaches `pause_target`; in-flight reasoning programs are marked to pause after completion. Pausing clears the program assignment. Normal resume selects a fairness-ordered prefix against aggregate capacity, then uses largest-first best-fit-decreasing packing across workers. A request that waits longer than `resume_timeout_seconds` is assigned to the worker with the most capacity after decayed acting-token usage and force-released.
+
+The request target is advisory rather than a reservation. Caller constraints and worker liveness remain authoritative, and `Sent` corrects the classifier's accounting if selection falls back.
+
+The classifier retains `ClassifyRequest::progress()` and reads its context high-water mark during reconciliation, so pressure decisions account for generation before completion. Dynamo initializes the counter from scheduling input tokens and raises it as the host observes prompt plus generated tokens; this is a logical token estimate, not physical KV occupancy. The handle survives migration retries, and completion records the terminal context size. A final session removes its program when the final request is admitted; completion or abort cannot restore it. A continuing idle program retains its observed assignment for `session_retention_seconds`. Idle retention is pruned lazily on the next classification or periodic reconciliation. The tracking limit also bounds retained programs; at the limit, the oldest idle retained program is evicted before admitting a new session.
 
 ## Configuration
 
-The policy registers the type `thunderagent`. Its YAML parameters are the fields of `ThunderAgentConfig`; [`worker-selection.yaml`](worker-selection.yaml) supplies the defaults.
+Link this crate in place of Dynamo's default router-plugin catalog, enable soft Dynamo session affinity, and select ThunderAgent for both classification and aggregated worker selection:
+
+```yaml
+request_classifier:
+  type: thunderagent
+  parameters:
+    pause_threshold: 0.95
+    pause_target: 0.80
+    resume_hysteresis: 0.10
+    resume_timeout_seconds: 1800
+    session_retention_seconds: 1800
+    scheduler_interval_seconds: 5
+    acting_token_weight: 1
+    acting_decay_tau_seconds: 1
+    buffer_per_program: 100
+    max_tracked_requests: 10000
+worker_selection:
+  aggregated: thunderagent
+  instances:
+    - name: thunderagent
+      type: thunderagent
+      parameters: {}
+```
 
 | Parameter | Default | Meaning |
 | --- | ---: | --- |
-| `pause_threshold` | `0.95` | Usage fraction that begins pausing programs |
-| `pause_target` | `0.80` | Usage fraction to recover to after pausing |
-| `resume_hysteresis` | `0.10` | Additional headroom required before greedy resume |
-| `resume_timeout_seconds` | `1800` | Maximum deferred duration before forced resume |
-| `session_retention_seconds` | `1800` | Retention period for completed acting sessions |
-| `scheduler_interval_seconds` | `5` | Capacity reconciliation cadence |
-| `acting_token_weight` | `1` | Normal capacity weight for an acting program |
-| `acting_decay_tau_seconds` | `1` | Half-life control for the decayed acting estimate |
-| `buffer_per_program` | `100` | Reserved tokens per active program |
-
-## Integrate as a policy catalog
-
-Link the crate directly as Dynamo's `dynamo-worker-selection-policy-catalog` dependency, or call [`register`](src/lib.rs) from a combined catalog. The policy factory can also be constructed programmatically:
-
-```rust
-use std::sync::Arc;
-
-use dynamo_kv_router::WorkerSelectionPolicyFactory;
-use thunderagent_dynamo_policy::{ThunderAgentConfig, worker_selection_policy};
-
-let config = ThunderAgentConfig::default();
-let factory: WorkerSelectionPolicyFactory = Arc::new(move |router, worker_type, _partition| {
-    worker_selection_policy(router.clone(), worker_type, config.clone())
-        .expect("valid ThunderAgent configuration")
-});
-```
+| `pause_threshold` | `0.95` | Worker usage fraction that begins pausing programs |
+| `pause_target` | `0.80` | Worker usage fraction that a pause cycle drains toward |
+| `resume_hysteresis` | `0.10` | Headroom below the pause threshold required for normal resume |
+| `resume_timeout_seconds` | `1800` | Maximum classifier deferral before forced release |
+| `session_retention_seconds` | `1800` | Idle continuing-program retention period |
+| `scheduler_interval_seconds` | `5` | Fixed cadence for pressure and resume decisions while programs are tracked |
+| `acting_token_weight` | `1` | Capacity weight for a program during tool work |
+| `acting_decay_tau_seconds` | `1` | Half-life in seconds for acting-token usage during timeout fallback placement |
+| `buffer_per_program` | `100` | Fixed token headroom per active program |
+| `max_tracked_requests` | `10000` | Independent bounds for active request state and retained program state before classification reaches Dynamo's queue |