Repository navigation
feat(llm): add native budget admission decorators - #43
Merged
senamakel merged 91 commits intoOct 21, 2026
Merged
Conversation
Adds a provider-neutral BeforeDispatch trait and a builder method on BudgetedModel so hosts can recheck an absolute deadline or admission policy after waiting for monetary capacity. The hook runs before each provider invocation or stream construction, and a refusal refunds the reservation without reaching the provider. Auto-committed-on: dragonfly
The budgeted model now consults its BeforeDispatch hook before each physical call and refunds the reservation when the hook refuses, so a refused dispatch never reaches the provider. The hook was previously only reachable through the builder but never actually invoked. Auto-committed-on: dragonfly
Reordered the public re-exports in the model module so budget comes before decorators, and reformatted the probe model's response construction in the budget tests to satisfy rustfmt. No behaviour changed. Auto-committed-on: dragonfly
The before-dispatch hook checks in the invoke and stream paths were rewritten as let-chain conditionals, flattening the nested if-let and if-let blocks without changing behaviour. Auto-committed-on: dragonfly
…spatch-admission Recheck host admission after monetary budget waits
Update the lockfile to rustls-webpki 0.103.16 and webpki-roots 1.0.10, and align the getrandom entry pulled in by tempfile to 0.3.4. Auto-committed-on: dragonfly
Add tests asserting that the before-dispatch hook is re-evaluated before each physical retry and when a lazy stream is first polled, so a revoked dispatch stops further calls. Auto-committed-on: dragonfly
The before-dispatch hook was only consulted once per request, so retries and fallback attempts could reach the provider without re-checking it. The hook is now stored on the attempt scope and invoked from before_physical_attempt, ensuring each physical call is gated. Auto-committed-on: dragonfly
The documentation for the BeforeDispatch hook and the before_dispatch installer now states that the check runs before each physical provider attempt, including retries and a lazy stream's first poll, rather than only once before dispatch. Auto-committed-on: dragonfly
Wrap the stream failure assertion across multiple lines so it fits the formatter's line width limit. No behaviour change. Auto-committed-on: dragonfly
Documented that implementations of BeforeDispatch must tolerate repeated checks for the same logical call, since the check runs at each physical transport attempt. A committed host permit should make the operation idempotent. Auto-committed-on: dragonfly
Adds a DispatchCheckedModel wrapper that applies a host-owned BeforeDispatch policy to any provider without requiring a monetary budget, checking before each physical send and once before generic model work for adapters lacking a separate send seam. This lets hosts enforce dispatch policy on unbudgeted providers that previously bypassed the check. Auto-committed-on: dragonfly
The retry and lazy stream tests now keep a handle on the budget and check its snapshot after dispatch, confirming that reservations are released and spend is recorded as expected. Auto-committed-on: dragonfly
Add tests asserting that when a before-dispatch hook refuses the first physical send, both unary and streaming OpenAI calls return the validation error, leave the budget snapshot untouched, and never reach the network. Auto-committed-on: dragonfly
A refused dispatch hook now records that the refusal happened before any physical attempt, so the reservation is only refunded when no request was actually sent. Previously a retry refused after the first physical send could refund a call that had already been charged. Auto-committed-on: dragonfly
Reformatted assertions and the dispatch hook in the budget tests so they match rustfmt output. No test behaviour changed. Auto-committed-on: dragonfly
The retry-refusal test now uses a hook that refuses only from the third dispatch attempt, so the first physical send is actually performed before the refusal. The local server waits for the invoke to finish and then asserts that no second connection arrives, confirming the charge for the first send is kept while the retry is refused. Auto-committed-on: dragonfly
The unbudgeted physical attempt path now records a refusal on the scope before propagating the error, so a nested dispatch refusal correctly refunds the outer budget reservation instead of leaving it charged. Auto-committed-on: dragonfly
The dispatch check in invoke and stream now records a refusal before returning the error, so a request rejected before any physical attempt is accounted for the same way as one refused inside a provider. The refusal-marking logic was also factored into a shared helper used by the existing physical attempt path. Auto-committed-on: dragonfly
Refusal tests now race the model call against the listener accept instead of waiting on a fixed timeout, so a connection opened by a refused dispatch fails the test immediately rather than only when it lands within the window. The shared assertion also checks the accept error kind, making the no-connection guarantee explicit. Auto-committed-on: dragonfly
Add a test asserting that when a budgeted model is wrapped in an outer dispatch check, both gates are consulted before the call, the first physical send, and the retry, and that a revoked outer gate blocks the retry. The existing no-connection assertion is factored into a shared helper. Auto-committed-on: dragonfly
Extend the nested dispatch refusal test to attach a dispatch-check hook to the budgeted model and assert it is invoked twice, confirming checks still run when an outer budget reservation is refunded. Auto-committed-on: dragonfly
Reformatted several long assertions in the budget tests to satisfy rustfmt line-length limits, and raised the local server accept timeout from one to five seconds so the retry test is less likely to flake on slow machines. Auto-committed-on: dragonfly
The nested hooks retry test now reads the incoming request and writes a valid chat completion response instead of dropping the connection, so the retry path is exercised against a real HTTP exchange. The no-connection assertion only runs when no retry occurred. Auto-committed-on: dragonfly
The test now discards the result of sending on the finished channel instead of unwrapping it, so a closed receiver no longer panics the test. Auto-committed-on: dragonfly
The physical attempt guard now runs the unbudgeted admission check even when a dispatch hook is present, so a wrapping DispatchCheckedModel's independent scope is enforced on every physical send including provider retries. The hook check was also collapsed into a let-chain. Auto-committed-on: dragonfly
The doc comment referenced `before_physical_attempt` as an intra-doc link, but no such item is in scope, so the link was replaced with plain inline code to avoid a broken rustdoc reference. Auto-committed-on: dragonfly
…lbacks Add tests asserting that an untracked paid adapter is still charged when a nested gate refuses after the send, and that an Anthropic fallback rechecks dispatch before issuing its second request instead of bypassing the gate. Auto-committed-on: dragonfly
Await the second server result before asserting on it and move the dispatch check ahead of the outcome match so the test fails on the bypass itself rather than on an unrelated mismatch. Auto-committed-on: dragonfly
Add a trait method reporting whether a model checks the scoped admission hook before every physical provider send. It defaults to false so callers cannot assume a refusal after invocation began happened before any provider request. Auto-committed-on: dragonfly
The reset of the untracked inner invocation flag now happens when a dispatch is refused before invocation rather than when a physical attempt is refused before attempt, so the flag is cleared at the point that actually corresponds to the invocation being skipped. Auto-committed-on: dragonfly
The test helper no longer clears the untracked invocation flag when no attempts have been recorded, since the flag is already false at that point and the reset was dead code. Auto-committed-on: dragonfly
The helper now returns just the scope it was given instead of also prepending any physical call scope from the task-local context. This removes a redundant lookup whose result was not needed by callers. Auto-committed-on: dragonfly
…l call scopes The DispatchCheckedModel no longer forwards fetch_deferred to its inner model, and nested_attempt_scopes now seeds its result from the current physical call scopes before appending the given scope. This keeps nested attempt tracking consistent with the ambient physical call context. Auto-committed-on: dragonfly
Add a fetch_deferred implementation to DispatchCheckedModel that delegates to the wrapped model, so deferred status lookups work through the dispatch-checking wrapper instead of being unavailable. Auto-committed-on: dragonfly
The dispatch check no longer requires the inner model to lack a pre-invocation admission gate before marking untracked inner invocation, since physical attempt tracking alone determines that condition. Auto-committed-on: dragonfly
The untracked-inner-invoked marker was only set when the inner model both lacked physical attempt tracking and had no pre-invocation admission gate. Removing the second condition ensures the marker is set whenever physical attempts are untracked, so models with an admission gate are no longer incorrectly treated as tracked. Auto-committed-on: dragonfly
The dispatch check now also skips marking the inner model as untracked when it has a pre-invocation admission gate, since such a gate already accounts for the invocation and marking it would misreport the attempt. Auto-committed-on: dragonfly
The budgeted model now leaves the untracked-inner flag alone when the inner model performs its own pre-invocation admission, since that path already accounts for the attempt and marking it would double count. Auto-committed-on: dragonfly
Collapse four wrapped boolean conditions in the dispatch and budget model wrappers onto one line each. Purely a formatting change with no behavioural difference. Auto-committed-on: dragonfly
AttemptScope now records whether the wrapped model has a pre-invocation admission gate, and can_refund_before_first_attempt treats that gate as equivalent to a refusal before the first attempt. This lets a reservation be refunded when the inner model rejects the call up front instead of leaving the budget charged for work that never ran. Auto-committed-on: dragonfly
Auto-committed-on: dragonfly
Deferred fetches previously bypassed the dispatch admission gate and attempt tracking that other calls go through. The fetch now runs the before-dispatch check, records untracked inner invocations, and executes inside a dispatch hook scope so deferred work is accounted for the same way as regular invocations. Auto-committed-on: dragonfly
The warmup budget tests now use a biased select so the completion signal is checked before accepting another connection, and they assert that no second connection arrives when the warmup finishes first. This makes the tests fail loudly instead of silently passing if an extra physical attempt slips through. Auto-committed-on: dragonfly
Adds new test cases to verify budget behavior when using models that track physical attempts, specifically testing scenarios involving nested budgeting and retries. Auto-committed-on: dragonfly
Auto-committed-on: dragonfly
Auto-committed-on: dragonfly
Auto-committed-on: dragonfly
Adds a new test case to verify that budget refunds are handled correctly when a gated fallback model refuses invocation before any child model is called. This ensures budget state remains unchanged in this specific failure scenario. Auto-committed-on: dragonfly
Introduces a method to check if all models in the chain have a pre-invocation admission gate. This allows for early rejection of requests before any model processing begins. Auto-committed-on: dragonfly
This removes the check for pre-invocation admission gates across all models in the fallback chain. This simplifies the fallback logic by removing an unnecessary pre-check. Auto-committed-on: dragonfly
Introduces a method to check if all models in the chain have a pre-invocation admission gate. This allows for early rejection of requests before any model processing begins. Auto-committed-on: dragonfly
This updates the assertion in budget tests to explicitly check that the fallback mechanism was not invoked when the primary dispatch refused before any child calls. This improves test clarity and robustness. Auto-committed-on: dragonfly
Adds a test case to verify that when a retry is refused due to live capacity, the budget is not marked as exhausted. This ensures that temporary capacity issues do not lead to incorrect budget state reporting. Auto-committed-on: dragonfly
This updates a test case to verify that when a retry is denied by live money reservations, it does not consume an additional host slot. This ensures correct behavior regarding resource accounting during retries. Auto-committed-on: dragonfly
Adds detailed documentation to explain that retry admission is synchronous and how budget exhaustion during a retry affects refusal reporting. This clarifies the behavior when a wait-mode budget cannot reserve capacity for a retry. Auto-committed-on: dragonfly
This updates the assertion in the budget tests to use explicit line breaks for better readability. This change ensures the test remains clear while verifying the budget cost matches the call policy cost. Auto-committed-on: dragonfly
…2-halt-owner-tinyinference-sync # Conflicts: # crates/tinyinference-llm/src/model/mod.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add native inference budget decorators and admission support so TinyAgents can reserve dispatches, apply shared budgets, preserve cumulative gateway charges, and expose the budget/error state needed by host model wrappers. The change includes native budget slots, fallback/decorator propagation, and TinyHumans model registration support.
Validation at
721380e728a55ee024ac26ef6952eb225e4e64ab:cargo check -p tinyinference-llm(pass)cargo clippy -p tinyinference-llm --all-targets -- -D warnings(pass)cargo test -p tinyinference-llm(838 tests plus all package doctests passed)cargo fmt --all --check(pass)Part of tinyhumansai/openhuman#8092.