Repository navigation
perf(ai-gateway): skip the group policy transaction for non-enterprise orgs - #8101
Closed
chrarnoldus wants to merge 1 commit into
Closed
chrarnoldus wants to merge 1 commit into
chrarnoldus wants to merge 1 commit into
Conversation
…e orgs Non-Enterprise organizations always evaluate to the unrestricted policy, which depends only on the organization's data_collection setting. The balance query already loads the plan and settings and joins the caller's membership, so the gateway now builds that policy directly instead of opening the read-only policy transaction (5 replica round trips) on every org request. The balance row now also reports organizations.deleted_at so soft-deleted organizations keep taking the policy-context path, which rejects them.
Contributor
Code Review SummaryStatus: No Issues Found | Recommendation: Merge Executive SummaryThe non-Enterprise policy shortcut matches Files Reviewed (8 files)
Reviewed by deepseek-v4.1-flash-1017 · Input: 57.4K · Output: 10.9K · Cached: 466.6K Review guidance: REVIEW.md from base branch |
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.
Summary
Stop opening the organization model-access policy transaction for organizations that are not on the Enterprise plan.
The transaction from
getOrganizationGroupPolicyContextruns on every organization request and is one of the most frequent gateway queries. Each run isbegin isolation level repeatable read read only, the membership select, theorganization_group_policy_settingsselect, the group memberships join, andcommit: 5 sequential round trips per org request. #8073 routes it throughreadDb, but it is still 5 round trips before inference on every org request.For non-Enterprise organizations,
evaluateEffectiveModelAccessPolicyalways returns the same unrestricted policy. It depends only onorganization.settings.data_collectionandcontext.policyRevision. The gateway already loads the plan and settings throughgetBalanceAndOrgSettings(…, readDb), so the transaction adds nothing on that path.nonEnterpriseModelAccessPolicy(dataCollection, policyRevision)inmodel-access.server.tsis the single source for that policy.evaluateEffectiveModelAccessPolicyreturns it for non-Enterprise contexts, so the gateway shortcut and the evaluated policy cannot drift.llm-proxy.ts,organizationGroupPolicyPromisewaits on the balance query and returnsnonEnterpriseModelAccessPolicy(settings?.data_collection, 0)when the plan is present, notenterprise, and the organization is not soft-deleted. Every other case keeps the existinggetOrganizationGroupPolicyContextcall.findOrganizationMemberBalanceRowalso selectsorganizations.deleted_at, andgetBalanceForOrganizationUserexposes it asorganizationDeleted. This keeps soft-deleted organizations failing as they do now (see Reviewer Notes).Verification
NODE_ENV=test pnpm exec jest -i --forceExit ../../packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts ../../packages/web-shared/src/lib/organizations/effective-model-access.server.test.ts ../../packages/web-shared/src/lib/organizations/organization-usage.test.ts ../../packages/web-shared/src/lib/ai-gateway/auto-routing-denied-models.test.ts src/lib/organizations/legacy-model-restrictions-parity.test.ts(fromapps/web, against this worktree's own Postgres fromKILO_PORT_OFFSET=auto pnpm test:db): 5 suites, 247 tests passedpnpm --filter auto-routing exec vitest run src/custom-router.test.ts: 38 passedpnpm --filter @kilocode/worker-utils exec vitest run src/account-access.test.ts: 7 passedpnpm --filter @kilocode/web-shared --filter web --filter ai-gateway typecheck: passedpnpm --filter auto-routing --filter @kilocode/worker-utils typecheck: passedpnpm exec oxlint --config .oxlintrc.json <8 changed files>: 0 warnings, 0 errorspnpm format:changed,git diff --check: cleanNew tests:
getOrganizationGroupPolicyContextorevaluateEffectiveModelAccessPolicy. Its access check getsnonEnterpriseModelAccessPolicy('deny', 0), and both an allowed decision (200) and a denied decision (404model_not_allowed) still apply.readDb, and their evaluated policy reaches the access decision.evaluateEffectiveModelAccessPolicyon a non-Enterprise context (with Enterprise-only settings and group policies present) strictly equalsnonEnterpriseModelAccessPolicy.getBalanceForOrganizationUserreportsorganizationDeletedbefore and after a soft delete (DB test).Visual Changes
N/A
Reviewer Notes
Membership is still enforced. When
subject.allowNonMemberis unset, the transaction throwsUNAUTHORIZEDunless the caller has a directorganization_membershipsrow. That check already happens twice before the shortcut:getUserFromAuth→validateUserAuthorization(packages/web-shared/src/lib/user/server.ts) callsisOrganizationMember(readDb, organizationId, kiloUserId)whenever the request carries an organization ID. It returns 403Access denied (not a member of the organization)otherwise. The shortcut only runs whenauth.authFailedResponseis null.findOrganizationMemberBalanceRow(packages/worker-utils/src/account-access.ts) inner joinsorganization_membershipson the caller (excludingbilling_manager).planis only set when that row exists. The shortcut requires a plan, so a missing membership row (or a billing manager) still takes the transaction path, exactly as before.All three reads (auth, balance, and the transaction since #8073) use
readDb, so the shortcut does not change replica-lag behavior either.Soft-deleted organizations. Neither the auth check nor the balance query filters
organizations.deleted_at, andmarkOrganizationAsDeletedkeeps membership rows. Today the policy context is the only gateway check that rejects a member of a soft-deleted organization, by throwingNOT_FOUND. Skipping it without a deleted check would have loosened access. The balance row now selectsdeleted_at, and a soft-deleted organization keeps taking the context path. This is additive: no existing balance consumer changes behavior. Theauto-routingtest fixture gainsdeleted_at: nullfor the widened row type.policyRevision. On
EffectiveOrganizationModelPolicy, nothing in production readspolicyRevision. Searching forpolicyRevisionand.policyRevisionfinds only the producers inmodel-access.server.ts, the policy-settings writers inorganization-groups.ts, and the unrelated sub-organization editor. None of the consumers on the gateway path read it:getEffectiveModelDecision,policyNeedsCandidateEvaluation,deniedAutoRoutingCandidates,deniedModelIdsForCandidates,autoRoutingDenials,resolveAccessCheck, andisAutoFreeCandidateAllowed. The shortcut uses0, which is what the context reports when the organization has noorganization_group_policy_settingsrow. A non-Enterprise org that was once Enterprise may still have a row with a higher revision. That difference is unobservable on this path.Consumers of
organizationGroupPolicyPromisebehave identically for non-Enterprise orgs:resolveAccessCheckstill callsgetEffectiveModelDecisionwith the same policy object. Model-dependent outcomes, such as restriction-exempt models or Kilo-exclusive provider restrictions, are evaluated as before.autoRoutingDenials: the settings deny-list fallback only applies when!groupPolicy && plan === 'enterprise', so it stays off.collectDeniedAutoRoutingModelIdsanddeniedAutoRoutingCandidatesreturn[]for an unrestricted grant, and privacy denials come fromeffectivePrivacy, not from the policy.isAutoFreeCandidateAllowedstill evaluatesgetEffectiveModelDecisionagainst the same policy.Settings shape. The balance row's
settingsandcontext.organization.settingsare the sameorganizations.settingsjsonb column with the sameOrganizationSettingstype. Neither side parses it at runtime.settingsis onlyundefinedon the no-user branch, which returnsnullbefore the shortcut.Latency trade-off. The policy promise now starts after the balance query instead of alongside it. Non-Enterprise orgs save 5 round trips. Enterprise orgs (and the fallback cases) start the transaction one region-local replica query later. That transaction is awaited much later in the handler (auto-routing or access check), so most of that query still overlaps other work.
Sibling PR overlap. A parallel PR adds a Redis cache for the Enterprise policy and touches the same
organizationGroupPolicyPromiselines. This diff is limited to the gate in front ofgetOrganizationGroupPolicyContext, so the cache should only need to wrap the remaining context branch when rebasing.