Add Cache-Aware "Token + Cache" Usage Limit Type - #2876
Draft
Kaveh-Vakili wants to merge 1 commit into
Draft
Conversation
Adds a new per-model, cache-token-aware usage limit alongside the existing Token/Compute time restrictions. Admins can configure a cache read/write weight (0-1000%) per model via Model Settings or the SMSS file, and members can be restricted with a new "Token + Cache" limit that counts cache read/creation tokens toward their token budget at that weighted percentage, on top of normal prompt+completion tokens. - Constants: add CACHE_READ_WEIGHT/CACHE_WRITE_WEIGHT metadata keys and MODEL_TOKEN_CACHE_RESTRICTION_VALUE restriction value - SecurityOwlCreator: declare CACHEREADWEIGHT/CACHEWRITEWEIGHT columns on MODELMETADATA - SecurityModelMetadataUtils: validate (0-1000%), persist, and seed the two new fields from smss properties - ModelUsageRestrictionUtility: new token_cache branch (engine- and user-level) that resolves a model's configured weights and folds weighted cache tokens into the usage check and running total - ModelInferenceLogsUtils: sum CACHE_READ_TOKENS/CACHE_CREATION_TOKENS alongside MESSAGE_TOKENS for token-family restriction modes - AbstractModelEngineResponse: carry resolved weights between the limit check and the per-call update Also fixes a pre-existing bug found while testing: getEngineUsagePermissionMap's SMSS_USER/ENGINEPERMISSION left outer join used bare table names and was silently resolving to zero rows for every restriction mode, not just this one. Switched to explicit qualified join columns, matching the existing pattern in SecurityProjectUtils. Note: MODELMETADATA is a pre-existing table; the OWL schema declaration above does not auto-ALTER existing databases, so any environment upgrading onto an existing security DB needs the two columns added manually (ALTER TABLE MODELMETADATA ADD COLUMN CACHEREADWEIGHT DOUBLE / CACHEWRITEWEIGHT DOUBLE) before this will work.
Contributor
✅ Snyk checks have passed. No issues have been found so far.
💻 Catch issues earlier using the plugins for VS Code, JetBrains IDEs, Visual Studio, and Eclipse. |
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.
Description
Adds a new per-model, cache-token-aware usage limit alongside the existing Token/Compute time restrictions. Admins can set a cache read/write weight (0–1000%) per model in Model Settings (or directly via the model's SMSS file), and members can be restricted with a new Token + Cache limit type that counts cache read/creation tokens toward their token budget at that weighted percentage, on top of normal prompt+completion tokens. Also fixes a pre-existing bug in the engine-permission lookup that was silently returning no restriction data at all, for any restriction mode — found while testing this feature live.
Changes Made
Constants.java: add CACHE_READ_WEIGHT, CACHE_WRITE_WEIGHT metadata keys and MODEL_TOKEN_CACHE_RESTRICTION_VALUE ("token_cache") restriction value
SecurityOwlCreator.java: declare CACHEREADWEIGHT/CACHEWRITEWEIGHT double columns on MODELMETADATA
SecurityModelMetadataUtils.java: wire the two new fields through validation (0–1000%, 0 allowed), the ModelMetadata record, insert/update SQL, row mapping, and smss-property seeding (upsertModelMetadata)
ModelUsageRestrictionUtility.java: new token_cache branch (engine- and user-level) that resolves a model's configured weights and folds weighted cache tokens into both the window-sum check and the per-response running total; plain token mode is unchanged (0% cache weight)
ModelInferenceLogsUtils.java: extend the usage-sum queries to also pull CACHE_READ_TOKENS/CACHE_CREATION_TOKENS and apply the caller-supplied weights for token-family modes
AbstractModelEngineResponse.java: add cacheReadWeight/cacheWriteWeight keys for passing resolved weights between the limit check and the per-call update
SecurityEngineUtils.java: fix — getEngineUsagePermissionMap's SMSS_USER↔ENGINEPERMISSION join used bare table names for a left.outer.join, which was silently resolving to zero rows regardless of restriction type; switched to explicit qualified join columns (SMSS_USER__ID / ENGINEPERMISSION__USERID), matching the pattern already used in SecurityProjectUtils.java
How to Test
Notes
Deployment gap, not covered by this PR: MODELMETADATA is a pre-existing table. SecurityOwlCreator's schema declaration only reconciles OWL/semantic metadata, not physical columns — it will not auto-ALTER TABLE an existing security DB to add the two new columns. Any environment upgrading onto an existing security DB needs to run manually first:
ALTER TABLE MODELMETADATA ADD COLUMN CACHEREADWEIGHT DOUBLE;
ALTER TABLE MODELMETADATA ADD COLUMN CACHEWRITEWEIGHT DOUBLE;
otherwise GetModelMetadata/UpdateModelMetadata will throw. Worth a migration step or startup check in a follow-up.
The SecurityEngineUtils join fix affects all restriction modes, not just token_cache — flagging in review since it's a behavior change to existing enforcement, found incidentally while testing this feature.