Skip to content

perf(ai-gateway): skip the group policy transaction for non-enterprise orgs - #8101

Closed
chrarnoldus wants to merge 1 commit into
mainfrom
perf/gateway-skip-non-enterprise-policy
Closed

chrarnoldus wants to merge 1 commit into
mainfrom
perf/gateway-skip-non-enterprise-policy

Conversation

@chrarnoldus

@chrarnoldus chrarnoldus commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Stop opening the organization model-access policy transaction for organizations that are not on the Enterprise plan.

The transaction from getOrganizationGroupPolicyContext runs on every organization request and is one of the most frequent gateway queries. Each run is begin isolation level repeatable read read only, the membership select, the organization_group_policy_settings select, the group memberships join, and commit: 5 sequential round trips per org request. #8073 routes it through readDb, but it is still 5 round trips before inference on every org request.

For non-Enterprise organizations, evaluateEffectiveModelAccessPolicy always returns the same unrestricted policy. It depends only on organization.settings.data_collection and context.policyRevision. The gateway already loads the plan and settings through getBalanceAndOrgSettings(…, readDb), so the transaction adds nothing on that path.

  • nonEnterpriseModelAccessPolicy(dataCollection, policyRevision) in model-access.server.ts is the single source for that policy. evaluateEffectiveModelAccessPolicy returns it for non-Enterprise contexts, so the gateway shortcut and the evaluated policy cannot drift.
  • In llm-proxy.ts, organizationGroupPolicyPromise waits on the balance query and returns nonEnterpriseModelAccessPolicy(settings?.data_collection, 0) when the plan is present, not enterprise, and the organization is not soft-deleted. Every other case keeps the existing getOrganizationGroupPolicyContext call.
  • findOrganizationMemberBalanceRow also selects organizations.deleted_at, and getBalanceForOrganizationUser exposes it as organizationDeleted. 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 (from apps/web, against this worktree's own Postgres from KILO_PORT_OFFSET=auto pnpm test:db): 5 suites, 247 tests passed
  • Mutation checks on the new llm-proxy tests: removing the soft-delete guard fails 2 tests, and disabling the shortcut fails 2 tests
  • pnpm --filter auto-routing exec vitest run src/custom-router.test.ts: 38 passed
  • pnpm --filter @kilocode/worker-utils exec vitest run src/account-access.test.ts: 7 passed
  • pnpm --filter @kilocode/web-shared --filter web --filter ai-gateway typecheck: passed
  • pnpm --filter auto-routing --filter @kilocode/worker-utils typecheck: passed
  • pnpm exec oxlint --config .oxlintrc.json <8 changed files>: 0 warnings, 0 errors
  • pnpm format:changed, git diff --check: clean

New tests:

  • A non-Enterprise org does not call getOrganizationGroupPolicyContext or evaluateEffectiveModelAccessPolicy. Its access check gets nonEnterpriseModelAccessPolicy('deny', 0), and both an allowed decision (200) and a denied decision (404 model_not_allowed) still apply.
  • Enterprise orgs, members without a spending membership (no plan from the balance query), and soft-deleted orgs still call the policy context with readDb, and their evaluated policy reaches the access decision.
  • A rejection from the policy context for a soft-deleted org still fails the request before upstream.
  • evaluateEffectiveModelAccessPolicy on a non-Enterprise context (with Enterprise-only settings and group policies present) strictly equals nonEnterpriseModelAccessPolicy.
  • getBalanceForOrganizationUser reports organizationDeleted before and after a soft delete (DB test).

Visual Changes

N/A

Reviewer Notes

Membership is still enforced. When subject.allowNonMember is unset, the transaction throws UNAUTHORIZED unless the caller has a direct organization_memberships row. That check already happens twice before the shortcut:

  1. getUserFromAuth → validateUserAuthorization (packages/web-shared/src/lib/user/server.ts) calls isOrganizationMember(readDb, organizationId, kiloUserId) whenever the request carries an organization ID. It returns 403 Access denied (not a member of the organization) otherwise. The shortcut only runs when auth.authFailedResponse is null.
  2. findOrganizationMemberBalanceRow (packages/worker-utils/src/account-access.ts) inner joins organization_memberships on the caller (excluding billing_manager). plan is 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, and markOrganizationAsDeleted keeps membership rows. Today the policy context is the only gateway check that rejects a member of a soft-deleted organization, by throwing NOT_FOUND. Skipping it without a deleted check would have loosened access. The balance row now selects deleted_at, and a soft-deleted organization keeps taking the context path. This is additive: no existing balance consumer changes behavior. The auto-routing test fixture gains deleted_at: null for the widened row type.

policyRevision. On EffectiveOrganizationModelPolicy, nothing in production reads policyRevision. Searching for policyRevision and .policyRevision finds only the producers in model-access.server.ts, the policy-settings writers in organization-groups.ts, and the unrelated sub-organization editor. None of the consumers on the gateway path read it: getEffectiveModelDecision, policyNeedsCandidateEvaluation, deniedAutoRoutingCandidates, deniedModelIdsForCandidates, autoRoutingDenials, resolveAccessCheck, and isAutoFreeCandidateAllowed. The shortcut uses 0, which is what the context reports when the organization has no organization_group_policy_settings row. 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 organizationGroupPolicyPromise behave identically for non-Enterprise orgs:

  • resolveAccessCheck still calls getEffectiveModelDecision with 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. collectDeniedAutoRoutingModelIds and deniedAutoRoutingCandidates return [] for an unrestricted grant, and privacy denials come from effectivePrivacy, not from the policy.
  • isAutoFreeCandidateAllowed still evaluates getEffectiveModelDecision against the same policy.

Settings shape. The balance row's settings and context.organization.settings are the same organizations.settings jsonb column with the same OrganizationSettings type. Neither side parses it at runtime. settings is only undefined on the no-user branch, which returns null before 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 organizationGroupPolicyPromise lines. This diff is limited to the gate in front of getOrganizationGroupPolicyContext, so the cache should only need to wrap the remaining context branch when rebasing.

…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.
@chrarnoldus chrarnoldus self-assigned this Oct 19, 2026
@kilo-code-bot

kilo-code-bot Bot commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Code Review Summary

Status: No Issues Found | Recommendation: Merge

Executive Summary

The non-Enterprise policy shortcut matches evaluateEffectiveModelAccessPolicy field-for-field, membership and soft-delete behavior is preserved, and the added deleted_at field is correctly selected and propagated, with no blocking issues found.

Files Reviewed (8 files)
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts
  • 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/group-policies/model-access/model-access.server.ts
  • packages/web-shared/src/lib/organizations/organization-usage.ts
  • packages/web-shared/src/lib/organizations/organization-usage.test.ts
  • packages/worker-utils/src/account-access.ts
  • services/auto-routing/src/custom-router.test.ts

Reviewed by deepseek-v4.1-flash-1017 · Input: 57.4K · Output: 10.9K · Cached: 466.6K

Review guidance: REVIEW.md from base branch main

@chrarnoldus
chrarnoldus deleted the perf/gateway-skip-non-enterprise-policy branch October 19, 2026 14:20
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