feat(ext): add route selector extension point and adaptive strategy - #611
feat(ext): add route selector extension point and adaptive strategy#611SantiagoDePolonia wants to merge 3 commits into
Conversation
Extensions can now steer load balancing for virtual models. A new ext.RouteSelector registers through the registry's single selector slot and is consulted by redirects using the new "adaptive" strategy; it picks among the currently viable targets (catalog-supported, with rate-limit capacity) and observes every upstream attempt — primaries, retries, and failovers — through the llmclient hooks, so it can score targets by real traffic. Session affinity, capacity filtering, failover chains, and retries all remain core's responsibility. Without a registered selector — or when the selector declines, answers outside the pool, or panics (contained and logged) — the adaptive strategy behaves exactly like round_robin, so configs stay portable between core and extended builds. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Warning Review limit reached
Next review available in: 17 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds the ChangesAdaptive routing
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ExtensionRegistry
participant App
participant ProviderFactory
participant VirtualModelService
participant RouteSelector
ExtensionRegistry->>App: expose registered selector
App->>ProviderFactory: install attempt hooks
App->>VirtualModelService: install selector
VirtualModelService->>RouteSelector: select from RouteRequest
ProviderFactory->>RouteSelector: report attempt lifecycle
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@config/config.example.yaml`:
- Line 64: Update the strategy comment near the strategy configuration example
to describe the routing extension as optional rather than required, while
preserving that adaptive is valid without a selector and falls back to
round_robin.
In `@ext/registry_test.go`:
- Around line 88-102: Convert TestRegistryRouteSelectorSingleSlot and
TestRouteTargetQualified into table-driven tests covering unset, replacement,
and nil-reset selector behavior plus route-target qualification. Reuse each
case’s inputs and expected outputs through subtests, preserving the current
assertions and behavior.
In `@internal/app/app.go`:
- Around line 117-127: Update the recovery handler inside routeSelectorHooks so
it performs no extension calls, including selector.Name(), and logs only fixed,
non-sensitive metadata such as the observation event. Remove the recovered panic
value from slog.Error while preserving panic containment and the existing error
message.
In `@internal/virtualmodels/adaptive_test.go`:
- Around line 82-88: Expand the request-translation assertions in the relevant
test to compare the complete req.Candidates slice against expected
ext.RouteCandidate metadata, including Weight, OutputPerMtok, InputPerMtok,
provider/model fields, and nil versus populated optional prices. Convert the
coverage to table-driven tests so multiple candidate and parameter-mapping cases
are validated, while preserving the existing request translation behavior.
In `@internal/virtualmodels/adaptive.go`:
- Around line 20-25: Update the deferred recovery handler around Select so panic
reporting does not call the extension-provided selector.Name(); use a constant
label or panic-safe helper while preserving the resolvedTarget{} false fallback.
Add a regression test covering both Select and Name panicking and verify
round-robin fallback is returned without another panic.
- Around line 22-23: Update the panic recovery logging around the route selector
fallback to stop including the raw recovered value r. Remove the panic field or
replace it with a non-sensitive panic classification, while preserving the
selector, source, and fallback behavior.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: fbb3af64-6418-4421-a897-87bf4dde574b
📒 Files selected for processing (14)
config/config.example.yamlconfig/virtualmodels.goext/ext.goext/registry.goext/registry_test.goext/route.gointernal/admin/handler_virtualmodels.gointernal/app/app.gointernal/virtualmodels/adaptive.gointernal/virtualmodels/adaptive_test.gointernal/virtualmodels/balancer.gointernal/virtualmodels/service.gointernal/virtualmodels/types.gointernal/virtualmodels/validation.go
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if r := recover(); r != nil { | ||
| slog.Error("route selector panicked; falling back to round robin", | ||
| "selector", selector.Name(), "source", entry.vm.Source, "panic", r) |
There was a problem hiding this comment.
When a selector callback and its Name method both panic, the recovery handler calls selector.Name() while processing the first panic, allowing the second panic to escape and abort adaptive resolution instead of falling back to round robin. The observation recovery in internal/app/app.go has the same failure mode and can abort an upstream request.
| // An extension route selector observes every upstream attempt — primaries, | ||
| // retries, and failovers — to steer adaptive load balancing. Like the | ||
| // health tracker, its hooks must be attached before any provider exists. | ||
| var routeSelector ext.RouteSelector | ||
| if cfg.Extensions != nil { | ||
| routeSelector = cfg.Extensions.RouteSelector() | ||
| } | ||
| if routeSelector != nil { | ||
| cfg.Factory.AddHooks(routeSelectorHooks(routeSelector)) |
There was a problem hiding this comment.
When an upstream request fails on retryable attempts before succeeding, these request-level hooks emit only one final successful outcome because the transport retry loop runs inside a single llmclient request scope. This breaks the public per-attempt observation contract and hides intermediate failures from adaptive selectors, causing unstable targets to receive inaccurately healthy scores.
Artifacts
Repro: focused mocked-upstream Go test and RouteSelector callback recorder
- Evidence file captured while the check ran.
Repro: verbose failing test output showing three HTTP attempts but only two selector callbacks
- The full command output behind this check.
| if model, found := s.catalog.LookupModel(t.qualified); found && model != nil && model.Metadata != nil && model.Metadata.Pricing != nil { | ||
| candidate.InputPerMtok = model.Metadata.Pricing.InputPerMtok | ||
| candidate.OutputPerMtok = model.Metadata.Pricing.OutputPerMtok |
There was a problem hiding this comment.
These fields expose the catalog's mutable *float64 values directly to extension code because LookupModel performs only a shallow model copy. A selector that modifies either value changes prices subsequently read by cost routing and cost calculation, and concurrent mutation can also introduce a data race; pass defensive copies instead.
Confidence Score: 2/5The change should not merge until selector panics are fully contained and retry attempts are reported according to the public contract. Recovery can itself panic through selector.Name(), while transport retries are collapsed into one observation despite the API promising one observation per attempt; candidate pricing also exposes mutable shared catalog values. Files Needing Attention: internal/app/app.go, internal/virtualmodels/adaptive.go, ext/route.go
T-Rex LogsWhat T-Rex did
|
Panic recovery around selector calls no longer re-enters extension code: the selector's name is captured panic-safe at install time, so a Name that panics mid-recovery can no longer escape the handler and abort the request instead of falling back. The recovered value is no longer logged — it is extension-controlled and could carry request data. Candidate pricing is now defensively copied: selectors received the catalog's own *float64 values, so a misbehaving extension could mutate the prices that cost routing and cost calculation later read. The per-attempt observation contract is corrected to match behavior: hooks fire once per upstream call, and transport-level retries inside the provider client are aggregated into that call's single outcome (final status, duration spanning retry backoff). Moving observation inside the client retry loop would change hook semantics for every consumer (metrics, health tracking) and is out of scope here. Also: registry selector-slot tests are table-driven (incl. nil reset), the delegation test asserts the full candidate set, a regression test covers the double-panic path, and the example config no longer implies the adaptive strategy requires an extension. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Addressed the review findings in 189beca:
#612 has been rebased on top. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/virtualmodels/adaptive_test.go`:
- Around line 119-135: Expand TestBalancer_AdaptiveCandidatePricingIsCopied into
a table-driven isolation test that mutates both InputPerMtok and OutputPerMtok
for every selector candidate, then verifies the corresponding pricing pointers
for every catalog model remain unchanged. Use the existing candidate and catalog
model symbols to cover all candidates and both pricing fields, while preserving
the current defensive-copy assertion behavior.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 10507680-ad9f-47ea-a5ff-fc0f9baa0dc0
📒 Files selected for processing (7)
config/config.example.yamlext/registry_test.goext/route.gointernal/app/app.gointernal/virtualmodels/adaptive.gointernal/virtualmodels/adaptive_test.gointernal/virtualmodels/service.go
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Follow-up from the second review pass in b3c07a4: the pricing-isolation test now mutates every pricing pointer of every candidate and verifies all catalog prices stay intact. |
Summary
Adds a routing extension point to the public
extAPI so extension builds can steer load balancing for virtual models:ext.RouteSelector(ext/route.go):Selectpicks among currently viable targets of a load-balanced virtual model;OnAttemptStart/OnAttemptEndobserve every upstream attempt (primaries, retries, failovers) via the llmclient hooks, so selectors score targets from real traffic. Registered through a single slot on the registry (RegisterRouteSelector).adaptiveload-balancing strategy for virtual models: delegates target choice to the registered selector. Candidates carry provider/model, configured weight, and registry pricing.internal/app/app.go): selector hooks attach before providers are created (same rule as the health tracker); the virtual-models service gets the selector alongside the capacity probe.Safety / portability
round_robinbehavior. Configs usingadaptivestay portable between plain core and extended builds.Tests
ext: registry single-slot semantics,RouteTarget.Qualified().virtualmodels: selector delegation (request shape incl. pricing), fallback matrix (no selector / decline / out-of-pool answer / panic), single-target bypass.make test-race).🤖 Generated with Claude Code
Summary by CodeRabbit