Skip to content

Add Cache-Aware "Token + Cache" Usage Limit Type - #2876

Draft
Kaveh-Vakili wants to merge 1 commit into
devfrom
cache-token-weighted-limits
Draft

Add Cache-Aware "Token + Cache" Usage Limit Type#2876
Kaveh-Vakili wants to merge 1 commit into
devfrom
cache-token-weighted-limits

Conversation

@Kaveh-Vakili

Copy link
Copy Markdown

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

  1. On a model's Settings page, set Cache read weight / Cache write weight (e.g. 10 / 200) and save — confirm via GetModelMetadata(engine=[""]); that both persist.
  2. In that model's Access Control, set a member's Usage Limit Type to Token + Cache with a small Max Tokens value.
  3. Run GetUserModelUsageRestrictions(engine=[""]); as that member — should return restrictedBy: "token_cache", currentValue, maxValue, and the resolved cacheReadWeight/cacheWriteWeight (default 100 if unset on the model).
  4. Regression-check plain Token mode still returns correctly from the same reactor (validates the getEngineUsagePermissionMap fix didn't regress the existing path).
  5. Send real chat traffic against the model until over the limit — confirm the IllegalArgumentException fires with the usual "Token limit exceeded..." message.

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.

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.
@snyk-io

snyk-io Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Snyk checks have passed. No issues have been found so far.

Status Scan Engine Critical High Medium Low Total (0)
Open Source Security 0 0 0 0 0 issues

💻 Catch issues earlier using the plugins for VS Code, JetBrains IDEs, Visual Studio, and Eclipse.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant