Skip to content

[AISAGE-707] refactor(ai-gateway): resolve custom LLMs by catalog lookup and fail closed - #7780

Merged
chrarnoldus merged 8 commits into
mainfrom
kilo/private-custom-llm-flag
Oct 20, 2026
Merged

chrarnoldus merged 8 commits into
mainfrom
kilo/private-custom-llm-flag

Conversation

@chrarnoldus

@chrarnoldus chrarnoldus commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Code no longer decides that a model is a custom LLM by checking whether its id starts with kilo-internal/. It asks a custom LLM catalog instead, and that prefix is no longer required. Access to custom LLMs is also hardened ahead of the free/public behaviour in #7778 and #7291.

Custom LLM catalog (custom-llm/custom-llm-catalog.ts)

  • Each instance caches the parsed custom_llm2 rows, keyed by lowercased public id. Lookups use findCustomLlm and isCustomLlm.
  • The cache expires after 60 s. There is no cross-instance invalidation, so admin edits apply everywhere within a minute.
  • After a failed read, the catalog waits 5 s before querying the database again.
  • Two accessors:
    • getCustomLlmsById returns an empty map when nothing could be read.
    • requireCustomLlmsById throws instead. Model stats publication uses it, so an unreadable catalog never exposes custom LLM scores.

Routing fails closed

  • getProvider checks the catalog first, before direct BYOK, ChatGPT and Vercel BYOK.
  • A custom LLM id resolves to custom-llm-unavailable when either:
    • the caller has no access: anonymous, a personal account, or an organization or group that is not allowlisted
    • the stored credentials are missing or invalid
  • It never falls back to Vercel or OpenRouter.
  • The proxy answers custom-llm-unavailable with the same 404 as an unknown model id, so callers cannot probe for custom LLMs.

Model policy exemption requires a grant

  • isModelRestrictionExempt now takes a CustomLlmGrantee (an organization and the member's groups).
  • A custom LLM is exempt only when that organization or one of those groups is in its organization_ids / group_ids. Otherwise organization and group model policy applies as for any model.
  • EffectiveOrganizationModelPolicy carries the grantee from the policy context. The legacy ModelRestrictions path takes it as an optional field, and without one no custom LLM is exempt.

Public id rules (admin editor and router)

  • The kilo-internal/ prefix is no longer required, and CUSTOM_LLM_PREFIX is removed.
  • A new id (created or a copy target) must match the auto free model id format (GatewayModelIdSchema).
  • It must not contain : or ~, or start with kilo/, kilo-auto/, kilocode/ or a direct BYOK provider prefix.
  • It must not be an OpenRouter or Kilo-exclusive model, or differ only in case from an existing row.
  • The editor validates the format before saving.
  • Existing rows keep their ids and can still be edited.

Other prefix checks replaced

  • Organization Auto: a target that is a custom LLM is rejected.
  • Model stats: terminal-bench summaries skip custom LLMs, and Enkrypt sync and publication withhold their scores.
  • Catalog flag: entries from listAvailableCustomLlms carry isCustomLlm: true. Auto-routing pool validation, the custom router dialog, the admin benchmark options, the default model dialog and the mode picker use the flag instead of the prefix.

This is a prep PR for #7291.

Verification

  • No manual testing.
  • Covered by tests for:
    • catalog lookup, cache expiry, failure backoff and the strict accessor
    • fail-closed routing for other organizations, personal accounts, anonymous callers and missing credentials
    • the proxy's 404
    • grant-scoped policy exemption by organization and by group
    • every id rule
    • model stats and Enkrypt withholding

Visual Changes

The admin Custom LLMs editor accepts ids without kilo-internal/ and shows the id rules. I took no screenshots.

Reviewer Notes

  • On a cold instance the first lookup loads all custom_llm2 rows from the replica.
  • Model stats tests forbid replica reads, so they stub the catalog.

@chrarnoldus chrarnoldus self-assigned this Oct 16, 2026
@kilo-code-bot

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

Copy link
Copy Markdown
Contributor

Code Review Summary

Status: No Issues Found | Recommendation: Merge

Executive Summary

The incremental commit only stubs the custom LLM catalog in the openrouter-models test suite, and the two stubbed exports cover every catalog symbol reachable from the routes under test.

Overview

Severity Count
CRITICAL 0
WARNING 0
SUGGESTION 0
Files Reviewed (1 file)
  • apps/web/src/tests/openrouter-models.test.ts - 0 issues

Verification Notes

  • Incremental review of 8e5271d0..d5b1fd5e; the only change is a jest.mock for custom-llm/custom-llm-catalog adding getCustomLlmsById and requireCustomLlmsById returning empty maps.
  • Verified the mock satisfies the test's import graph: model-stats-cache.ts:7 and terminal-bench.ts:1 import requireCustomLlmsById (terminal-bench is also separately mocked), and the mocked organization-models module cuts off the isCustomLlm path. findCustomLlm/isCustomLlm are only used by get-provider.ts and model-allow.server.ts, neither of which is reachable from the models, transcription-models, or stats routes exercised here.
  • The factory only closes over new Map(), so it cannot trigger Jest's out-of-scope variable error, and each call returns a fresh empty map, which preserves the previous "no custom LLMs" behavior of the suite.
  • Previous WARNING on custom-llm-catalog.ts remains resolved; no catalog code changed in this increment. No memory-leak risk introduced (no listeners, timers, or retained references added).
  • No markdown documentation files changed, so the image-format rule does not apply.
  • Tests were not executed (read-only review); this assessment is from source inspection only.
Previous Review Summaries (4 snapshots, latest commit 8e5271d)

Current summary above is authoritative. Previous snapshots are kept for context only.

Previous review (commit 8e5271d)

Status: No Issues Found | Recommendation: Merge

Executive Summary

The incremental commit makes every benchmark/Enkrypt publication path fail closed when the custom LLM catalog has never been read, and the changed lines introduce no new defects.

Overview

Severity Count
CRITICAL 0
WARNING 0
SUGGESTION 0
Files Reviewed (8 files)
  • apps/web/src/lib/model-stats/model-stats-cache.test.ts - 0 issues
  • apps/web/src/lib/model-stats/sync-enkrypt.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/custom-llm/custom-llm-catalog.test.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/custom-llm/custom-llm-catalog.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/enkrypt.test.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/model-stats-cache.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/terminal-bench.ts - 0 issues
  • packages/web-shared/src/lib/organizations/organization-models.test.ts - 0 issues

Verification Notes

  • Incremental review of 522abb32..8e5271d0; the earlier WARNING is resolved. requireCustomLlmsById throws only when no catalog value has ever been loaded, and the model-stats snapshot, terminal-bench, and Enkrypt sync all call it instead of the lenient getCustomLlmsById. Routing still uses the lenient lookup, where a stale/empty result cannot publish scores.
  • Fail-closed behavior verified by reading the current code: getModelStatsSnapshot rethrows on a cold catalog failure and serves the previous snapshot otherwise; createCachedFetch in terminal-bench falls back to the last summaries or an empty map (omitting scores); syncEnkryptBenchmarks aborts and records the failure.
  • No memory leak introduced: the catalog keeps one cache entry and one in-flight promise, released in finally.
  • No markdown documentation files changed, so the image-format rule does not apply.
  • Tests were not executed (read-only review); this assessment is from source inspection only.

Previous review (commit 522abb3)

Status: 1 Issue Found | Recommendation: Address before merge

Executive Summary

Replacing the static kilo-internal/ check with a cached catalog makes benchmark/Enkrypt publication fail open when the catalog read fails, so custom LLM scores can be exposed.

Overview

Severity Count
CRITICAL 0
WARNING 1
SUGGESTION 0
Issue Details (click to expand)

WARNING

File Line Issue
packages/web-shared/src/lib/ai-gateway/custom-llm/custom-llm-catalog.ts 83 getCustomLlmsById returns an empty map after a failed load with no cache; the publication and Enkrypt sync paths treat that as "no custom LLMs", so replica-read failures can publish scores for protected models. The model-stats snapshot also caches the empty set for 5 minutes.
Files Reviewed (34 files)
  • apps/web/src/app/admin/auto-routing/eligible-model-options.ts - 0 issues
  • apps/web/src/app/admin/auto-routing/eligible-model-options.test.ts - 0 issues
  • apps/web/src/app/api/models/stats/[slug]/route.test.ts - 0 issues
  • apps/web/src/app/api/models/stats/route.test.ts - 0 issues
  • apps/web/src/components/auto-routing/CustomRouterDialog.tsx - 0 issues
  • apps/web/src/components/organizations/custom-modes/ModeForm.tsx - 0 issues
  • apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx - 0 issues
  • apps/web/src/lib/auto-routing-pool-validation.ts - 0 issues
  • apps/web/src/lib/model-stats/enkrypt-identity.ts - 0 issues
  • apps/web/src/lib/model-stats/enkrypt-identity.test.ts - 0 issues
  • apps/web/src/lib/model-stats/sync-enkrypt.ts - 0 issues
  • apps/web/src/lib/model-stats/sync-enkrypt.test.ts - 0 issues
  • apps/web/src/lib/model-stats/model-stats-cache.test.ts - 0 issues
  • apps/web/src/routers/organizations/organization-settings-router.test.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/custom-llm/custom-llm-catalog.ts - 1 issue
  • packages/web-shared/src/lib/ai-gateway/custom-llm/custom-llm-catalog.test.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/custom-llm/listAvailableCustomLlms.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/model-utils.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/providers/get-provider.ts - 0 issues
  • packages/web-shared/src/lib/ai-gateway/providers/get-provider.custom-llm.test.ts - 0 issues
  • packages/web-shared/src/lib/model-allow.server.ts - 0 issues
  • packages/web-shared/src/lib/model-allow.server.test.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/enkrypt-publication.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/enkrypt-publication.test.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/enkrypt.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/enkrypt.test.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/model-stats-cache.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/terminal-bench.ts - 0 issues
  • packages/web-shared/src/lib/model-stats/terminal-bench.test.ts - 0 issues
  • packages/web-shared/src/lib/organizations/effective-model-access.server.test.ts - 0 issues
  • packages/web-shared/src/lib/organizations/organization-auto-model.ts - 0 issues
  • packages/web-shared/src/lib/organizations/organization-models.test.ts - 0 issues
  • packages/web-shared/src/lib/organizations/organization-types.ts - 0 issues
  • packages/web-shared/src/tests/helpers/custom-llm.helper.ts - 0 issues

Verification Notes

  • No memory leak introduced: the catalog cache is bounded by custom_llm2 row count, and inFlight is released in finally.
  • Reviewed at HEAD 522abb32; the previously reviewed commit 5e10cc99 is not an ancestor of HEAD, so this was a full PR review with fresh comment state.
  • The prior inline note on organization-types.ts (rename to isCustomLlm) is resolved: the current diff adds isCustomLlm.

Fix these issues in Kilo Cloud

Previous review (commit 5e10cc9)

Status: No Issues Found | Recommendation: Merge

Executive Summary

The incremental change since the previous review only relocates the existing isPrivateCustomLlm: true property within the same convert() object literal, a pure reordering with no behavioral effect and no new memory-leak risk.

Files Reviewed (1 file)
  • packages/web-shared/src/lib/ai-gateway/custom-llm/listAvailableCustomLlms.ts - 0 issues

Verification Notes

  • Incremental diff since 0cdea4dd moves isPrivateCustomLlm: true from the end of the returned object to directly after name; both properties remain present exactly once, so the produced catalog entry is unchanged.
  • No other lines changed in this increment; the previously reviewed behavior across the 8 PR files is unaffected.
  • No new listeners, subscriptions, timers, or retained references introduced.

Previous review (commit 0cdea4d)

Status: No Issues Found | Recommendation: Merge

Executive Summary

Reviewed the private-custom-LLM flag refactor across 8 files; verified that isPrivateCustomLlm is stamped only on entries from listAvailableCustomLlms and survives to every consumer (org catalog via getAvailableModelsForOrganization → tRPC z.custom output and OpenRouterModelsResponseSchema parse; non-org /api/v1/models never contains custom LLMs), so behavior is preserved with high confidence.

Files Reviewed (8 files)

  • apps/web/src/components/auto-routing/AutoRoutingModeCard.tsx
  • apps/web/src/components/auto-routing/AutoRoutingModeCard.test.ts
  • apps/web/src/components/organizations/custom-modes/ModeForm.tsx
  • apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx
  • apps/web/src/lib/auto-routing-pool-validation.ts
  • apps/web/src/lib/auto-routing-pool-validation.test.ts
  • packages/web-shared/src/lib/ai-gateway/custom-llm/listAvailableCustomLlms.ts
  • packages/web-shared/src/lib/organizations/organization-types.ts

Verification Notes

  • isPrivateCustomLlm: true is set only in listAvailableCustomLlms.convert (the sole custom-LLM catalog source), so prefix-based and flag-based filtering are equivalent for current data.
  • OpenRouterModelSchema gains the field, so the client OpenRouterModelsResponseSchema parse no longer strips it; the org model tRPC procedure uses z.custom<OpenRouterModelsResponse>(), which does not strip.
  • Non-org handleModelsRequest and useModelSelectorList(undefined) paths return only enhanced managed models, so no custom-LLM exclusion is lost there.
  • No new listeners, subscriptions, or timers were introduced; no memory-leak risk observed.
  • Prefix checks intentionally retained elsewhere (routing, model policy, Org Auto target, model stats) match the PR description.

Reviewed by deepseek-v4.1-flash-1017 · Input: 47.3K · Output: 5.1K · Cached: 442.9K

Review guidance: REVIEW.md from base branch main

@chrarnoldus
chrarnoldus force-pushed the kilo/private-custom-llm-flag branch from 5e10cc9 to e262c2d Compare October 19, 2026 12:35
@chrarnoldus chrarnoldus changed the title refactor(web): flag private custom LLMs in the model catalog instead of matching their prefix refactor(ai-gateway): recognize custom LLMs by catalog lookup instead of the kilo-internal/ prefix Oct 19, 2026
@chrarnoldus chrarnoldus changed the title refactor(ai-gateway): recognize custom LLMs by catalog lookup instead of the kilo-internal/ prefix refactor(ai-gateway): resolve custom LLMs by catalog lookup and fail closed Oct 19, 2026
@chrarnoldus chrarnoldus changed the title refactor(ai-gateway): resolve custom LLMs by catalog lookup and fail closed [AISAGE-707] refactor(ai-gateway): resolve custom LLMs by catalog lookup and fail closed Oct 20, 2026
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.

2 participants