feat: add iii-llm-router worker - #32
Conversation
Policy-based LLM routing brain. Unopinionated by design —
router ships with ZERO built-in model names, zero hardcoded
pricing, zero provider assumptions. You register whatever
models your gateway (LiteLLM, Bifrost, OpenRouter, vLLM,
custom) supports at runtime. The router never invents a model
name.
Functions (18):
- router::decide hot path
- router::policy_{create,update,delete,list,test}
- router::classify prompt heuristic only
- router::classifier_config category → model map (user sets)
- router::ab_{create,record,report,conclude}
- router::health_{update,list}
- router::model_{register,unregister,list}
- router::stats
Triggers: 18 HTTP (POST + GET).
Decide pipeline:
match policies → pick highest priority
A/B variant sampling (weighted) if running
"auto" policy → classifier → user-mapped model
unhealthy primary → fallback
over-budget → search registered models for cheaper fit
meeting min_quality (else keep original + flag)
no policy + classifier → classify → map
no policy + no classifier → empty model + reason
State (scope llm-router):
policies, ab_tests, ab_events, routing_log,
model_health, classifier, models
Tests: 17 passing — matching, priority, A/B sampling, classifier
mapping, auto-without-classifier, fallback, budget downgrade
with/without registered models, health thresholds, heuristic.
SDK: iii-sdk 0.11.0 stable
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 11 minutes and 11 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (12)
📝 WalkthroughWalkthroughA new Rust service Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant Router Handler
participant State Store
participant Classifier
participant Decision Engine
Client->>Router Handler: POST /decide<br/>(prompt, tenant, budget...)
activate Router Handler
Router Handler->>State Store: Load policies, A/B tests,<br/>model health, registrations
activate State Store
State Store-->>Router Handler: Data loaded
deactivate State Store
Router Handler->>Decision Engine: decide(request, policies,<br/>tests, health, classifier_cfg)
activate Decision Engine
Decision Engine->>Decision Engine: Match policies by<br/>tenant/feature/tags
alt Policy matched
Decision Engine->>Decision Engine: Select highest priority
alt action.model == "auto"
Decision Engine->>Classifier: heuristic_complexity(prompt)
activate Classifier
Classifier-->>Decision Engine: category, confidence
deactivate Classifier
Decision Engine->>Decision Engine: Map category to<br/>model via config
end
Decision Engine->>Decision Engine: Check model health<br/>& error_rate threshold
Decision Engine->>Decision Engine: Apply budget<br/>downgrade if needed
else No policy matched
Decision Engine->>Decision Engine: Try classifier path<br/>or A/B test selection
end
Decision Engine-->>Router Handler: RoutingDecision<br/>(model, reason, confidence)
deactivate Decision Engine
Router Handler->>State Store: Persist audit log entry
activate State Store
State Store-->>Router Handler: Logged
deactivate State Store
Router Handler-->>Client: { model, reason,<br/>confidence, policy_id }
deactivate Router Handler
sequenceDiagram
actor Admin
participant API
participant Policy Handler
participant State Store
Admin->>API: POST /policies/create<br/>(name, match_rule, action, ...)
activate API
API->>Policy Handler: JSON payload
activate Policy Handler
Policy Handler->>Policy Handler: Generate id if missing<br/>pol-{uuid}
Policy Handler->>Policy Handler: Set created_at_ms
Policy Handler->>State Store: state_set(scope, key,<br/>serialized Policy)
activate State Store
State Store-->>Policy Handler: OK
deactivate State Store
Policy Handler-->>API: { policy_id, created: true }
deactivate Policy Handler
API-->>Admin: 200 OK + policy_id
deactivate API
Note over Admin,State Store: Later: routing requests<br/>use this policy to decide
Admin->>API: POST /policies/test<br/>(prompt, tenant, ...)
activate API
API->>Policy Handler: Test request
activate Policy Handler
Policy Handler->>State Store: Load all policies,<br/>AB tests, health, registrations
activate State Store
State Store-->>Policy Handler: All data
deactivate State Store
Policy Handler->>Policy Handler: Filter policies,<br/>run decide() with seed 0
Policy Handler-->>API: { matched_policies: [...],<br/>decision: {...} }
deactivate Policy Handler
API-->>Admin: 200 OK + matched list & decision
deactivate API
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 20
🧹 Nitpick comments (1)
llm-router/src/manifest.rs (1)
8-27: Consider deriving the manifest from the same function metadata used for registration.The manifest duplicates the function IDs/descriptions from
main.rs, and the tests only check the count plus two IDs. A sharedconstlist would prevent future route/manifest drift.Also applies to: 35-49
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@llm-router/src/manifest.rs` around lines 8 - 27, The manifest in manifest.rs duplicates the function IDs/descriptions that are also defined when registering handlers in main.rs, causing potential drift; extract the canonical list of functions (IDs and descriptions) into a shared constant (e.g., a pub const FUNCTIONS: &[(&str,&str)] or a serde-serializable static structure) and have both main.rs registration code (the handler registration logic) and manifest.rs derive the manifest from that single shared symbol instead of hardcoding; update any tests that assert counts/IDs to reference the shared constant as the source of truth.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@llm-router/SPEC.md`:
- Around line 123-125: The spec incorrectly describes a model-backed "shadow
classifier" and extra latency; update SPEC.md to reflect that the shipped
handler uses a local heuristic classifier (not an LLM), so there is no
Haiku/Flash-Lite dependency, no ~200ms added latency, and the runtime does not
call out to a model. Replace mentions of "model-backed shadow classifier",
"Classifier prompt is stored in state and hot-swappable", and any
latency/dependency claims with a clear statement that the handler uses a
deterministic local heuristic classifier (name it if present in code, e.g., the
"local heuristic classifier" or the handler function/class that implements it),
that no external model or prompt storage is used at runtime, and that any future
model-backed variant would be optional and explicitly documented.
- Around line 37-70: The SPEC.md function catalog is out of sync with main.rs:
it lists router::decide_batch, router::health_check, and router::ab_evaluate
(cron triggers) but those routes aren't registered, and it omits the
model-management endpoints actually exposed by main.rs; update SPEC.md so the
published API matches the registered functions in main.rs by (1) adding entries
for any registered model-management endpoints found in main.rs (use their exact
handler names), (2) removing or marking
router::decide_batch/router::health_check/router::ab_evaluate as
deprecated/unimplemented if they are not registered, or add clear notes that
they are TODO and not currently exposed, and (3) ensure the cron-trigger list
includes only the handlers wired in main.rs; reference the handler names in
main.rs to ensure exact naming consistency (e.g., router::decide_batch,
router::health_check, router::ab_evaluate and the model-management handler
names) so clients can rely on the spec.
In `@llm-router/src/config.rs`:
- Around line 44-48: After deserializing RouterConfig in load_config, perform
explicit validation on the resulting cfg: check numeric ranges (e.g., ensure
health_skip_threshold_error_rate is within 0.0..=1.0 and other numeric
thresholds are sane) and ensure identifier/string fields on RouterConfig are
non-empty; if any validation fails return an Err (e.g., via anyhow::bail or
with_context) describing the invalid field so startup is aborted with a clear
message. Ensure this validation runs immediately after serde_yaml::from_str(...)
and before returning Ok(cfg) so invalid YAML values like -1 or 2 are rejected.
In `@llm-router/src/functions/ab.rs`:
- Around line 63-98: Lookup the test definition before creating/persisting
AbEvent: fetch the test record using the same state scope (use the test_id and
the existing state access pattern) and verify the provided variant_model exists
in the test's configured variants; if not, return an error instead of calling
state::state_set. Also validate numeric fields on the payload before persisting:
ensure quality_score is present or defaulted but within 0.0..=1.0, cost_usd is
>= 0.0, and latency_ms is non-negative; if any validation fails return an
IIIError::Handler (or appropriate error) and skip calling state::state_set for
the AbEvent (the AbEvent, key_event and state::state_set are the relevant
symbols to modify).
- Around line 201-228: Require and validate the winner_model: change the handler
to error if payload.get("winner_model") is missing, then verify the provided
winner string is one of the variants in the loaded AbTest (e.g., check against
AbTest.variants or the field that lists candidate model ids) and return a
handler error if not found. After marking test.status = "concluded" and saving
it with state::state_set(&iii, &cfg.state_scope, &key_test(&test_id), ...), also
update the matching policy to roll out the winning model (e.g., update whatever
state key holds the policy using state::state_set or call the existing
policy-update helper with the winning model and test metadata) so the spec
requirement is fulfilled; if you intentionally do not perform a rollout, return
a response field that clearly states no policy change occurred. Ensure error
messages reference test_id and winner_model for clarity.
In `@llm-router/src/functions/classify.rs`:
- Around line 24-35: After extracting the prompt string, trim it and reject
empty/whitespace-only values before any classification work: replace the raw
prompt usage with let prompt = payload.get("prompt").and_then(|v|
v.as_str()).ok_or_else(|| IIIError::Handler("missing
'prompt'".into()))?.to_string(); then immediately do let prompt =
prompt.trim().to_string(); and if prompt.is_empty() return
Err(IIIError::Handler("empty 'prompt'".into())); before calling
heuristic_complexity(&prompt) or proceeding with classifier_id resolution. This
ensures classify.rs rejects blank prompts consistently.
In `@llm-router/src/functions/decide.rs`:
- Around line 133-147: The load_typed function currently swallows JSON parse
failures and drops malformed entries; update it to log or surface parse errors
(at minimum log) including the state key/prefix and the serde error so malformed
policy/model/health records are visible. Specifically, in load_typed (params:
iii, cfg, prefix) where you iterate results from state::state_list and inspect
it.get("value"), replace the silent if let Ok(parsed) branch so that the Err
from serde_json::from_value is captured and logged (include prefix, the item
key/id from it if present, the raw value, and the parse error); optionally
aggregate/return an IIIError if the call site requires failing fast. Ensure logs
use the project’s existing logger (or return an error wrapped as IIIError)
rather than dropping the record.
- Around line 58-67: The RNG is currently seeded from now_ms() which causes
identical seeds for burst requests; change the RNG construction in this block so
it uses entropy-backed seeding instead of the millisecond timestamp—replace the
seed/seed_from_u64 usage that creates mut rng (the StdRng instance) with
StdRng::from_entropy() (rand 0.8) so the decide(&req, ctx, &cfg, &mut rng) call
receives a non-deterministic RNG; update any variable names (seed, rng)
accordingly and remove the now_ms() seed variable.
In `@llm-router/src/functions/health.rs`:
- Around line 27-39: The current extraction of the model ID via
payload.get("model").and_then(|v| v.as_str()) allows empty or whitespace-only
strings; reject those before writing state by trimming and validating non-empty.
After obtaining the model string (the variable named model used with
key(&model)), check model.trim().is_empty() and return an
IIIError::Handler("missing 'model'") (or similar) when blank; keep using
ModelHealth deserialization, set h.last_checked_ms, and call
state::state_set(&iii, &cfg.state_scope, &key(&model), ...) only with the
validated non-blank model ID. Ensure any subsequent uses (key, state_set) use
the trimmed/validated model variable.
In `@llm-router/src/functions/model.rs`:
- Around line 27-37: The ModelRegistration allows whitespace-only model IDs and
negative pricing which can corrupt routing/pricing; before setting
m.registered_at_ms and calling state::state_set (in the block using
ModelRegistration, m.model, decide::now_ms, key(&m.model)), validate that
m.model.trim().is_empty() and reject it with Err(IIIError::Handler("missing
'model'".into())) if so, and also validate that any price fields on
ModelRegistration (e.g., m.input_price, m.output_price) are present and
non-negative, returning Err(IIIError::Handler("invalid price" or similar)) when
negative; perform these checks immediately after deserializing into
ModelRegistration and before persisting with state::state_set.
In `@llm-router/src/functions/policy.rs`:
- Line 68: merge_policy currently silently ignores invalid JSON in rule
match/action fields and allows priority to wrap when casting i64→i32 (causing
wrong/no updates); update merge_policy (& any similar updater used later) to
validate incoming payload.match and payload.action JSON by attempting full
deserialization and returning an error if invalid, and validate payload.priority
is within i32 bounds before casting (reject out-of-range values), so invalid
fields cause the update to fail rather than be truncated/ignored; ensure the
function returns/propagates a clear error for these validation failures instead
of proceeding to mutate p.
- Around line 152-173: test_handler currently constructs DecideContext with only
policies (DecideContext { policies: &policies, ..DecideContext::default() })
which leaves classifier mappings, experiments, model health and registered
models empty and causes dry-run decisions (decide) to differ from production;
either populate the same full state used by the production decide path or
restrict this endpoint to only return policy matches. Fix by loading and passing
the full routing state (classifier mappings, ab_tests/experiments, model health,
registered models, budgets, etc.) into DecideContext before calling decide(&req,
ctx, &cfg, &mut rng), mirroring the state retrieval logic used by the real
request handler, or change the endpoint to return only the matched policies and
avoid calling decide to prevent misleading full decisions. Ensure you update
references around match_policy, DecideContext, and decide to use the richer
context or remove decide invocation accordingly.
In `@llm-router/src/functions/stats.rs`:
- Around line 36-46: The current stats handler in stats.rs calls
state::state_list("routing_log:") and then filters in-memory by
horizon/tenant/feature, which loads the entire audit log; change this to avoid
full scans by querying state with time-bucketed or range-prefixed keys (use a
timestamp-bucketed prefix when writing routing_log entries) or implement
paginated/scan-with-cursor reads from state::state_list so you only iterate
buckets since horizon, or replace this ad-hoc scan with pre-aggregated counters
per-day/per-tenant/per-feature that the router updates on write; update the code
paths around decide::now_ms(), horizon, and the loop over items to use the new
prefix/range/pagination or counters instead of loading all items into memory.
In `@llm-router/src/main.rs`:
- Around line 72-75: register_functions currently discards errors from
register_function_with and lets main() log "registered 18 functions..." even if
some failed; change register_functions to aggregate results from each
register_function_with call (e.g., collect a Vec<anyhow::Error> or return
Result<(), Vec<Error>>), return that to main(), and in main() check the returned
Result and fail startup (log errors and exit/non-ok return) if any registration
failed—mirror how register_triggers handles Err(e) so startup only logs
readiness when all register_function_with calls succeeded.
In `@llm-router/src/router.rs`:
- Around line 121-222: The decide function currently evaluates A/B tests early
and returns before applying health and budget checks; reorder the decision flow
so it follows SPEC.md (health → most-specific policy → classifier → A/B →
budget): keep matched = ctx.policies... and sort_by as-is, then (1) select the
most-specific policy (matched.first()) and resolve its model (handling "auto"
via heuristic_complexity and classifier.thresholds), (2) run
skip_unavailable(&chosen, ctx.health, cfg.health_skip_threshold_error_rate) and
if unhealthy apply policy.action.fallback, (3) apply budget constraint via
downgrade_to_fit(remaining, req, ctx.models) and only then return the
policy-based RoutingDecision; (4) if no policy matched, use classifier
(heuristic_complexity + ctx.classifier.thresholds) to pick a model and then
check health and budget for that model before returning; (5) only after
policy/classifier paths fail, evaluate ctx.ab_tests and pick_ab_variant, then
subject the AB-chosen variant to the same health and budget checks (and fallback
handling) before returning; update references to policy.action.fallback,
skip_unavailable, downgrade_to_fit, heuristic_complexity, ctx.ab_tests,
pick_ab_variant, and ctx.classifier.thresholds accordingly.
- Around line 56-67: The sum of weights in pick_ab_variant can overflow when
using u32; change accumulation to u64 by computing let total: u64 =
variants.iter().map(|v| v.weight as u64).sum(), use a u64 pick (let mut pick =
rng.gen_range(0..total)), and when comparing/subtracting cast each v.weight to
u64 (if pick < v.weight as u64 { ... } pick -= v.weight as u64;). Keep the early
return when total == 0 and otherwise return the chosen v.model.clone()
unchanged.
In `@llm-router/src/state.rs`:
- Around line 14-19: The current Ok(val) branch in state response parsing (used
by state::get and the similar state::list handling at the other block) swallows
malformed responses by returning None when "value" is missing or when "items"
isn't an array; change this to validate the response shape and return a
descriptive Err instead of Ok(None) on malformed backend responses. Concretely,
in the function handling the Ok(val) branch inspect val.get("value") (and in the
list path inspect val.get("items")) and if the key is missing or has the wrong
type produce an error (with context like "malformed state response: missing or
invalid 'value'/'items'") rather than filtering to None; keep the existing
null-checks for legitimate null values but do not convert missing/typed-mismatch
into empty results.
- Around line 6-11: All state TriggerRequest calls leave timeout_ms as None
causing potential hangs in hot-path helpers state_get, state_set, state_delete,
and state_list used by router::decide; set an explicit per-request timeout_ms
(e.g., a value under the 2000ms SLO after verifying iii-sdk timeout semantics)
on each TriggerRequest construction so the trigger will error quickly on a
stalled backend, and ensure any callers handle the timeout error path
consistently (update the TriggerRequest instances in the functions named
state_get/state_set/state_delete/state_list and any related error handling in
router::decide).
In `@llm-router/src/types.rs`:
- Around line 3-37: Add serde strictness to reject unknown JSON fields and
support legacy field name: annotate PolicyMatch, PolicyAction, Policy, and
AbTest with #[serde(deny_unknown_fields)] so unknown/misspelled control-plane
fields are rejected; update the Policy (and AbTest) match field definition
(currently pub match_rule: PolicyMatch with #[serde(default, rename = "match")])
to also accept the legacy "match_rule" by adding alias = "match_rule" in the
serde attribute (i.e., #[serde(default, rename = "match", alias =
"match_rule")]); keep existing serde(default) behavior for other fields and
apply deny_unknown_fields to each struct to prevent silent acceptance of typos.
---
Nitpick comments:
In `@llm-router/src/manifest.rs`:
- Around line 8-27: The manifest in manifest.rs duplicates the function
IDs/descriptions that are also defined when registering handlers in main.rs,
causing potential drift; extract the canonical list of functions (IDs and
descriptions) into a shared constant (e.g., a pub const FUNCTIONS:
&[(&str,&str)] or a serde-serializable static structure) and have both main.rs
registration code (the handler registration logic) and manifest.rs derive the
manifest from that single shared symbol instead of hardcoding; update any tests
that assert counts/IDs to reference the shared constant as the source of truth.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: bc9878ff-5305-4d9a-a392-996cf5e78f58
📒 Files selected for processing (19)
llm-router/Cargo.tomlllm-router/README.mdllm-router/SPEC.mdllm-router/build.rsllm-router/config.yamlllm-router/src/config.rsllm-router/src/functions/ab.rsllm-router/src/functions/classify.rsllm-router/src/functions/decide.rsllm-router/src/functions/health.rsllm-router/src/functions/mod.rsllm-router/src/functions/model.rsllm-router/src/functions/policy.rsllm-router/src/functions/stats.rsllm-router/src/main.rsllm-router/src/manifest.rsllm-router/src/router.rsllm-router/src/state.rsllm-router/src/types.rs
| ## Functions (18) | ||
|
|
||
| ### Core Routing | ||
|
|
||
| ``` | ||
| router::decide | ||
| Input: { | ||
| tenant, feature?, user?, | ||
| prompt: string, # first 500 chars is enough for classification | ||
| tags?: ["payments", "p0"], # application-level tags | ||
| budget_remaining_usd?: 1.50, # from budget::check | ||
| latency_slo_ms?: 2000, # max acceptable latency | ||
| min_quality?: "high"|"medium"|"low" | ||
| } | ||
| Output: { | ||
| model: "claude-haiku-4.5", | ||
| reason: "policy:support-default matched, budget constraint applied", | ||
| policy_id?: "pol-xxx", | ||
| ab_test_id?: "ab-xxx", | ||
| fallback: "gpt-4.1-mini", | ||
| confidence: 0.92 | ||
| } | ||
| Notes: The hot path. Evaluation order: | ||
| 1. Check model health (skip unavailable providers) | ||
| 2. Match policies by tenant → feature → tags (most specific wins) | ||
| 3. If policy says "auto", run complexity classifier | ||
| 4. If A/B test active for this scope, apply variant weights | ||
| 5. If budget constraint, cap to cheapest model that meets min_quality | ||
| 6. Return model + reason + fallback | ||
|
|
||
| router::decide_batch | ||
| Input: {requests: [{tenant, feature, prompt, ...}]} | ||
| Output: {decisions: [{model, reason, ...}]} | ||
| Notes: Batch version for pipeline workloads. Single round-trip. |
There was a problem hiding this comment.
Align the published API catalog with the registered functions.
The spec advertises router::decide_batch and cron-triggered router::health_check/router::ab_evaluate, but main.rs does not register those. It also omits the registered model-management endpoints from this trigger list. Please update the function/trigger catalog so clients do not integrate against missing routes.
Also applies to: 208-231
🧰 Tools
🪛 markdownlint-cli2 (0.22.0)
[warning] 41-41: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@llm-router/SPEC.md` around lines 37 - 70, The SPEC.md function catalog is out
of sync with main.rs: it lists router::decide_batch, router::health_check, and
router::ab_evaluate (cron triggers) but those routes aren't registered, and it
omits the model-management endpoints actually exposed by main.rs; update SPEC.md
so the published API matches the registered functions in main.rs by (1) adding
entries for any registered model-management endpoints found in main.rs (use
their exact handler names), (2) removing or marking
router::decide_batch/router::health_check/router::ab_evaluate as
deprecated/unimplemented if they are not registered, or add clear notes that
they are TODO and not currently exposed, and (3) ensure the cron-trigger list
includes only the handlers wired in main.rs; reference the handler names in
main.rs to ensure exact naming consistency (e.g., router::decide_batch,
router::health_check, router::ab_evaluate and the model-management handler
names) so clients can rely on the spec.
| Notes: Uses a cheap model (Haiku/Flash-Lite) as a shadow classifier. | ||
| Adds ~200ms but saves $$ by avoiding frontier models for simple queries. | ||
| Classifier prompt is stored in state and hot-swappable. |
There was a problem hiding this comment.
Correct the classifier behavior description.
These lines describe a model-backed shadow classifier, but the shipped handler uses the local heuristic classifier and does not call an LLM. This misstates latency, dependencies, and runtime behavior.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@llm-router/SPEC.md` around lines 123 - 125, The spec incorrectly describes a
model-backed "shadow classifier" and extra latency; update SPEC.md to reflect
that the shipped handler uses a local heuristic classifier (not an LLM), so
there is no Haiku/Flash-Lite dependency, no ~200ms added latency, and the
runtime does not call out to a model. Replace mentions of "model-backed shadow
classifier", "Classifier prompt is stored in state and hot-swappable", and any
latency/dependency claims with a clear statement that the handler uses a
deterministic local heuristic classifier (name it if present in code, e.g., the
"local heuristic classifier" or the handler function/class that implements it),
that no external model or prompt storage is used at runtime, and that any future
model-backed variant would be optional and explicitly documented.
State layer (state.rs):
- Add explicit 1.5s timeout_ms on every state::* trigger so a stalled
backend errors fast on the decide hot path instead of hanging.
- Handle three engine response envelopes for state::list:
{ "items": [...] } (0.11.0), bare array (0.11.2), null. Reject any
other shape with a descriptive error instead of silently returning
empty.
- state::get validates the response too; null values read as absent,
malformed shapes surface as errors.
types.rs:
- deny_unknown_fields on Policy, PolicyMatch, PolicyAction, AbTest —
typo'd control-plane fields now fail loudly instead of being silently
ignored.
- Policy/AbTest accept both "match" and legacy "match_rule" via serde
alias.
router.rs:
- Reorder decide() to policy → classifier → A/B, with health + budget
applied consistently on every path. A/B is the last-resort fallback
when no policy and no classifier match, instead of intercepting
before health/budget checks.
- pick_ab_variant accumulates weights as u64 so N variants with max-u32
weights can't overflow.
functions/decide.rs:
- Replace now_ms()-seeded StdRng with StdRng::from_entropy(); millisecond
collisions were biasing A/B variant picks on burst requests.
- load_typed logs parse failures with scope, prefix, key, and target
type so malformed state entries aren't silently dropped.
- Trim incoming prompts and reject empty/whitespace-only.
functions/classify.rs:
- Trim prompt; reject empty/whitespace.
functions/ab.rs:
- record_handler verifies variant_model belongs to the test before
persisting, and validates quality_score ∈ [0, 1], latency_ms ≥ 0,
cost_usd ≥ 0.
- conclude_handler requires winner_model and verifies it's one of the
test variants. Returns rollout_applied:false with a note that the
caller drives rollout via router::policy_update.
functions/health.rs:
- Trim + reject blank model ids. Validate error_rate ∈ [0, 1].
functions/model.rs:
- Trim + reject blank model ids in register and unregister.
- Reject negative/NaN input_per_1m / output_per_1m.
functions/policy.rs:
- merge_policy returns Result; invalid match/action JSON or
out-of-range priority now fail the update instead of being silently
truncated.
- test_handler loads the full DecideContext (policies + ab_tests +
health + classifier + models) via the same path as the production
decide handler, so dry-runs can't silently diverge.
- load_list helper parallels decide::load_typed, tolerates both
envelope shapes, logs parse failures.
functions/stats.rs:
- Narrow state::list prefix to the shared leading digits of
[horizon, now] so scans don't pull the entire audit log when a
recent window is requested. Hard cap of 50,000 entries per call,
with truncated:true surfaced in the response.
- Tolerate both envelope shapes; log malformed entries.
config.rs:
- Post-deserialize validation: state_scope and classifier_default_id
must be non-empty; stats_default_days ≥ 1; health_skip_threshold_
error_rate ∈ [0, 1]. Invalid YAML now aborts startup with a clear
message.
main.rs + manifest.rs:
- Extract the canonical FUNCTIONS list as a const in manifest.rs.
main.rs registration and build_manifest both derive from it, so the
published manifest and the registered handlers can't drift. Manifest
count test now references the const.
- Note: iii_sdk::III::register_function_with returns an infallible
FunctionRef, so error aggregation at that call site isn't possible.
register_triggers still handles Err cases per existing pattern.
Cargo.toml:
- Pin iii-sdk to exact 0.11.2 (no range).
Tests: 30 passing. cargo clippy -- -D warnings clean.
There was a problem hiding this comment.
Actionable comments posted: 10
♻️ Duplicate comments (1)
llm-router/src/functions/stats.rs (1)
51-52:⚠️ Potential issue | 🟠 MajorThe hard cap is applied after the full state scan is already materialized.
state::state_listreturns aVec, so broad prefixes can still load all matching routing logs into memory beforeSCAN_HARD_CAPstops processing. Use day buckets, pagination/cursors, or pre-aggregated counters so the storage read itself is bounded.Also applies to: 62-67
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@llm-router/src/functions/stats.rs` around lines 51 - 52, The code currently calls state::state_list(&iii, &cfg.state_scope, &prefix) which materializes a Vec and can load unbounded results before the SCAN_HARD_CAP is applied; change the read to a bounded/streaming approach — either iterate day-bucketed prefixes produced by scan_prefix(now, horizon) or swap to a paginated/streaming API (e.g., state::state_list_paginated or state::state_scan_cursor) that yields pages/cursors and stop after SCAN_HARD_CAP entries; update the loops that use the results (the blocks around the scan_prefix call and the second similar block at lines 62–67) to consume pages/streams and break when the hard cap is reached so storage reads are bounded.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@llm-router/src/config.rs`:
- Around line 5-18: The RouterConfig struct currently allows unknown YAML fields
which lets typos silently fall back to defaults; add the serde attribute to deny
unknown fields by annotating RouterConfig with #[serde(deny_unknown_fields)] so
serde will error on unexpected keys (e.g., misspelled
health_skip_threshold_errorrate) and force callers to fix config; update the
struct declaration where RouterConfig is defined to include this attribute
alongside the existing derive/serde attributes.
In `@llm-router/src/functions/decide.rs`:
- Around line 42-45: The code currently assigns classifier_id directly from
req.classifier_id, which allows all-whitespace strings to bypass the default;
update the assignment so you trim whitespace and treat empty results as None
before falling back to cfg.classifier_default_id. Specifically, transform
req.classifier_id (e.g., via map(|s| s.trim().to_string()).filter(|s|
!s.is_empty())) and then use unwrap_or_else(||
cfg.classifier_default_id.clone()) to set classifier_id, referencing the
existing classifier_id variable, req.classifier_id, and
cfg.classifier_default_id.
In `@llm-router/src/functions/model.rs`:
- Around line 96-103: The model_list code currently assumes each item from
state::state_list is an object with a "value" envelope, so change the filter_map
closure to also accept bare values: for each item, attempt to extract and
deserialize item.get("value") if present, otherwise attempt to deserialize the
entire item value itself into ModelRegistration (using serde_json::from_value on
the cloned Value). Update the closure used to build out: Vec<ModelRegistration>
so it tries both shapes (envelope first, fallback to bare) and only collects
successful deserializations; keep the rest of the call to state::state_list and
the ModelRegistration type intact.
In `@llm-router/src/functions/policy.rs`:
- Around line 32-37: After parsing or merging policies you must run semantic
validation to reject empty action.model and negative max_cost_per_request_usd;
add a new validator (e.g., validate_policy_semantics(&Policy) -> Result<()>) and
call it right after parse_policy in the create flow (where parse_policy(...)
returns p) and right after merge_policy in the patch flow (where merged policy
is produced), returning a client error if validation fails. Ensure the validator
checks that p.action.model is non-empty (or non-whitespace) and that
p.max_cost_per_request_usd is >= 0, and surface clear error messages from
validate_policy_semantics when calling state::state_set or the patch response
paths.
- Around line 124-131: The current policy_list parsing only extracts Policy
objects from envelope items with a "value" key, causing bare-value responses to
be dropped; update the closure that builds out (where items is mapped) to
attempt deserializing a Policy from it.get("value") first and, if that fails,
try deserializing directly from the item itself (e.g., attempt
serde_json::from_value(it.clone()) when no "value" exists or parsing the
envelope fails) so both envelope and bare state-list formats produce Policy
entries.
In `@llm-router/src/main.rs`:
- Around line 72-78: register_triggers currently swallows errors but main still
logs readiness; change register_triggers to return a Result<(), Error> (or
anyhow::Result) and update all callers (e.g., the main startup call that follows
register_functions(&iii, cfg.clone()) and the other call in the 133-162 region)
to propagate or ?-return that Result so startup fails when any HTTP trigger
registration fails; ensure register_triggers collects and returns the
first/error aggregate from its internal registration attempts and update call
sites to handle the Result instead of ignoring errors.
- Around line 46-59: The current Err branch for load_config(&cli.config)
silently falls back to RouterConfig::default() for any error; change it so
parse/validation errors abort instead. Update the match around
config::load_config(&cli.config) (which assigns router_config) to detect a
missing-file error (e.g., io::ErrorKind::NotFound or your config error variant)
and only then use RouterConfig::default() with a warning; for all other errors
(parse/validation) log the detailed error via tracing::error!(error = %e, path =
%cli.config, ...) and terminate startup (e.g., std::process::exit(1) or return
Err) so the program does not continue with wrong routing settings.
In `@llm-router/src/router.rs`:
- Around line 140-180: When returning fallback or a downgraded model, ensure you
re-run the same health check (skip_unavailable) before returning: after
computing fb = policy.action.fallback and after downgraded =
downgrade_to_fit(...), call skip_unavailable(&fb_or_downgraded, ctx.health,
cfg.health_skip_threshold_error_rate) and only return them if they are
considered available; if they are unavailable fall back to the next safe option
(e.g., try the other fallback, continue evaluation, or return the original
RoutingDecision with an appropriate reason and lower confidence). Update the
branches that currently return the fallback or downgraded RoutingDecision (the
blocks referencing policy.action.fallback, downgrade_to_fit, chosen, reason,
confidence) to perform this extra availability check before constructing the
RoutingDecision.
- Around line 123-124: Multiple policies can have the same priority causing
non-deterministic selection; add a deterministic tie-breaker by computing a
specificity metric for Policy (e.g., count presence of match_rule.tenant,
match_rule.feature, and number of match_rule.tags) and then sort matched first
by priority (descending), then by specificity (descending), then by a stable
unique key like p.id (ascending). Implement a helper like
policy_specificity(&Policy) -> usize that sums these components and update the
sort on matched (which is created from ctx.policies.iter().filter(|p|
match_policy(req, p))) to use a tuple comparator using (p.priority,
policy_specificity(p), p.id) with appropriate Reverse wrappers to preserve the
desired ordering.
In `@llm-router/src/types.rs`:
- Around line 49-68: The RoutingRequest and ModelRegistration structs silently
accept unknown JSON fields which can hide typos; add the attribute
#[serde(deny_unknown_fields)] to the RoutingRequest and ModelRegistration struct
declarations so serde will error on unexpected fields, update any tests or API
consumers that rely on permissive behavior, and run the test suite to ensure no
callers need migration; locate the structs by name (RoutingRequest,
ModelRegistration) in the types.rs diff and add the attribute immediately above
each struct definition.
---
Duplicate comments:
In `@llm-router/src/functions/stats.rs`:
- Around line 51-52: The code currently calls state::state_list(&iii,
&cfg.state_scope, &prefix) which materializes a Vec and can load unbounded
results before the SCAN_HARD_CAP is applied; change the read to a
bounded/streaming approach — either iterate day-bucketed prefixes produced by
scan_prefix(now, horizon) or swap to a paginated/streaming API (e.g.,
state::state_list_paginated or state::state_scan_cursor) that yields
pages/cursors and stop after SCAN_HARD_CAP entries; update the loops that use
the results (the blocks around the scan_prefix call and the second similar block
at lines 62–67) to consume pages/streams and break when the hard cap is reached
so storage reads are bounded.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d236cb0c-3faa-43b3-b1e6-164768c76c00
⛔ Files ignored due to path filters (1)
llm-router/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (14)
llm-router/Cargo.tomlllm-router/src/config.rsllm-router/src/functions/ab.rsllm-router/src/functions/classify.rsllm-router/src/functions/decide.rsllm-router/src/functions/health.rsllm-router/src/functions/model.rsllm-router/src/functions/policy.rsllm-router/src/functions/stats.rsllm-router/src/main.rsllm-router/src/manifest.rsllm-router/src/router.rsllm-router/src/state.rsllm-router/src/types.rs
✅ Files skipped from review due to trivial changes (1)
- llm-router/Cargo.toml
🚧 Files skipped from review as they are similar to previous changes (5)
- llm-router/src/manifest.rs
- llm-router/src/functions/classify.rs
- llm-router/src/functions/health.rs
- llm-router/src/state.rs
- llm-router/src/functions/ab.rs
…unds
config.rs:
- deny_unknown_fields on RouterConfig so misspelled keys error instead of
silently falling back to defaults.
types.rs:
- deny_unknown_fields on RoutingRequest and ModelRegistration.
functions/decide.rs:
- classifier_id: trim + treat whitespace-only as unset before falling back
to cfg.classifier_default_id.
- load_typed now delegates to state::parse_item, which tolerates both the
envelope and bare-value list shapes.
functions/policy.rs:
- validate_policy_semantics(): action.model must be non-empty,
max_cost_per_request_usd >= 0. Runs after create parse AND after
merge_policy so patch flows can't bypass.
- list_handler uses state::parse_item, accepts both envelope and bare
values (covers state::list shape drift).
- load_list helper reduced to a one-liner now that parse_item lives in
state.rs.
functions/model.rs + functions/health.rs + functions/ab.rs:
- list handlers and ab report event loader switched to state::parse_item
so bare-value responses aren't silently dropped.
functions/stats.rs:
- Swap inline envelope-handling to state::parse_item.
- Warn-log when state::list returns more items than SCAN_HARD_CAP so the
truncation is visible. Comment flags the SDK-side pagination gap.
main.rs:
- register_triggers returns Result and propagates errors to main, so a
failed HTTP trigger registration aborts startup instead of logging a
warning and continuing into "ready".
- load_config errors: NotFound → warn + default (unchanged behavior for
missing config). Any other error (parse failure, validation failure) →
error-log and return Err. Walks the anyhow source chain to avoid
brittle string matching.
router.rs:
- policy_specificity() helper: tenant(1) + feature(1) + |tags|.
- Sort matched policies by (priority desc, specificity desc, id asc) so
same-priority ties resolve deterministically instead of picking
whichever policy iteration happened to return first.
- Health re-check on fallback: if the policy fallback is also unhealthy,
return the primary with a lowered confidence and an explicit reason
rather than masking the degradation behind a bad fallback.
- Health re-check on budget-downgrade target: if the cheapest model that
fits the budget is unhealthy, still return it but with a lower
confidence and "(downgrade target unhealthy)" in reason.
state.rs:
- parse_item<T>() helper: envelope { value } first, bare-value second.
Shared across every caller so the fallback is consistent.
.gitignore + removed Cargo.lock from tracking (lockfiles don't belong in git).
Tests: 30 passing. cargo clippy -- -D warnings clean.
Policy-based LLM routing brain. Intentionally unopinionated — ships with zero built-in model names, zero hardcoded pricing, zero provider assumptions. Wraps any gateway (LiteLLM / Bifrost / OpenRouter / local vLLM / your own proxy) by sitting in front of it: gateway asks
router::decide→ router returns a model ID string → gateway forwards.Functions (18)
router::decide{model, reason, policy_id?, ab_test_id?, fallback?, confidence, request_id}router::policy_create/update/delete/list/testrouter::classify{complexity, confidence, suggested_model}(suggested_model respects your classifier map)router::classifier_config{id, thresholds: {simple/moderate/complex/expert → <your model id>}}router::ab_create/ab_record/ab_report/ab_concluderouter::health_update/health_listrouter::model_register/model_unregister/model_listrouter::stats18 HTTP triggers (POST + GET) under
/api/router/...Decide pipeline
State (scope
llm-router)What this is NOT
routing_logis audit-only; use iii's OTel for telemetry.router::classifier_configor wrapping a stronger classifier as a separate worker.Tests
17 passing: policy matching, priority ordering, A/B weighted sampling, classifier mapping, auto-without-classifier, unhealthy-fallback, budget-downgrade with and without registered models, health skip thresholds, heuristic category classification.
Stack
iii-sdk 0.11.0stablerandfor A/B weighted samplingserde_jsonfor all state blobsSummary by CodeRabbit