Skip to content

Move AI gateway state source of truth from Redis to PostgreSQL #5970

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chrarnoldus merged 3 commits into from
Sep 9, 2026
Merged

Move AI gateway state source of truth from Redis to PostgreSQL #5970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e7f1cae
refactor(ai-gateway): read state from database
chrarnoldus Sep 8, 2026
cea68b5
fix(ai-gateway): keep state mirrors consistent
chrarnoldus Sep 8, 2026
a35a775
fix(ai-gateway): retain stale alert TTL
chrarnoldus Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers-stale-alert.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ const FRESH_SYNC = new Date(NOW.getTime() - SYNC_PROVIDERS_STALE_AFTER_MS + 1);
const STALE_SYNC = new Date(NOW.getTime() - SYNC_PROVIDERS_STALE_AFTER_MS);
const OLDER_ALERT = new Date(STALE_SYNC.getTime() - 60_000);
const NEWER_ALERT = new Date(STALE_SYNC.getTime() + 60_000);
const EXPIRED_ALERT = new Date(NOW.getTime() - SYNC_PROVIDERS_STALE_ALERT_TTL_SECONDS * 1000);

afterEach(() => {
jest.restoreAllMocks();
Expand Down Expand Up @@ -109,6 +110,16 @@ describe('shouldPostStaleSyncAlert', () => {
).toBe(false);
});

it('alerts again when the last alert has expired', () => {
expect(
shouldPostStaleSyncAlert({
lastCompletedAt: null,
lastAlertAt: EXPIRED_ALERT,
now: NOW,
})
).toBe(true);
});

it('alerts again after a newer full sync goes stale', () => {
expect(
shouldPostStaleSyncAlert({
Expand Down Expand Up @@ -267,7 +278,7 @@ describe('alertIfSyncProvidersStale', () => {
expect(setLastAlertAt).not.toHaveBeenCalled();
});

it('swallows Redis failures so the cron can continue', async () => {
it('swallows state read failures so the cron can continue', async () => {
jest.spyOn(console, 'error').mockImplementation(() => undefined);
const sendNotification = jest.fn(
async (_notification: AdminSlackNotification) => 'posted' as const
Expand All @@ -277,7 +288,7 @@ describe('alertIfSyncProvidersStale', () => {
alertIfSyncProvidersStale({
now: () => NOW,
getLastCompletedAt: async () => {
throw new Error('redis down');
throw new Error('database down');
},
getLastAlertAt: async () => null,
sendNotification,
Expand Down
56 changes: 39 additions & 17 deletions apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers-stale-alert.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,16 @@ import { ai_gateway_sync_providers_state } from '@kilocode/db/schema';
import { APP_URL } from '@/lib/constants';
import { db } from '@/lib/drizzle';
import { redisClient } from '@/lib/redis';
import {
SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY,
SYNC_PROVIDERS_STALE_ALERT_LAST_POSTED_AT_REDIS_KEY,
} from '@/lib/redis-keys';
import { SYNC_PROVIDERS_STALE_ALERT_LAST_POSTED_AT_REDIS_KEY } from '@/lib/redis-keys';
import {
sendAdminSlackNotification,
type AdminSlackNotification,
} from '@/lib/slack/admin-notifications';
import { eq } from 'drizzle-orm';

export const SYNC_PROVIDERS_STALE_AFTER_MS = 60 * 60 * 1000;
export const SYNC_PROVIDERS_STALE_ALERT_TTL_SECONDS = 3 * 24 * 60 * 60;
const SYNC_PROVIDERS_STALE_ALERT_TTL_MS = SYNC_PROVIDERS_STALE_ALERT_TTL_SECONDS * 1000;

const STALE_WINDOW_LABEL = 'hour';
const STATUS_COPY = `No full sync has completed within the past ${STALE_WINDOW_LABEL}.`;
Expand Down Expand Up @@ -44,7 +43,11 @@ export function shouldPostStaleSyncAlert(input: {
) {
return false;
}
if (lastAlertAt !== null && (lastCompletedAt === null || lastAlertAt > lastCompletedAt)) {
if (
lastAlertAt !== null &&
(lastCompletedAt === null || lastAlertAt > lastCompletedAt) &&
now.getTime() - lastAlertAt.getTime() < SYNC_PROVIDERS_STALE_ALERT_TTL_MS
) {
return false;
}
return true;
Expand Down Expand Up @@ -167,25 +170,44 @@ type StaleAlertDependencies = {
};

async function defaultGetLastCompletedAt(): Promise<string | null> {
return redisClient.get<string>(SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY);
const [row] = await db
.select({ lastCompletedAt: ai_gateway_sync_providers_state.last_completed_at })
.from(ai_gateway_sync_providers_state)
.where(eq(ai_gateway_sync_providers_state.id, 1))
.limit(1);
return row?.lastCompletedAt ?? null;
}

async function defaultGetLastAlertAt(): Promise<string | null> {
return redisClient.get<string>(SYNC_PROVIDERS_STALE_ALERT_LAST_POSTED_AT_REDIS_KEY);
const [row] = await db
.select({ lastAlertAt: ai_gateway_sync_providers_state.stale_alert_last_posted_at })
.from(ai_gateway_sync_providers_state)
.where(eq(ai_gateway_sync_providers_state.id, 1))
.limit(1);
return row?.lastAlertAt ?? null;
}

async function defaultSetLastAlertAt(iso: string): Promise<unknown> {
const result = await redisClient.set(SYNC_PROVIDERS_STALE_ALERT_LAST_POSTED_AT_REDIS_KEY, iso, {
ex: SYNC_PROVIDERS_STALE_ALERT_TTL_SECONDS,
});
await db
.insert(ai_gateway_sync_providers_state)
.values({ stale_alert_last_posted_at: iso })
.onConflictDoUpdate({
target: ai_gateway_sync_providers_state.id,
set: { stale_alert_last_posted_at: iso },
return db.transaction(async tx => {
await tx.insert(ai_gateway_sync_providers_state).values({ id: 1 }).onConflictDoNothing();
const [row] = await tx
.select({ lastAlertAt: ai_gateway_sync_providers_state.stale_alert_last_posted_at })
.from(ai_gateway_sync_providers_state)
.where(eq(ai_gateway_sync_providers_state.id, 1))
.for('update');
if (!row) throw new Error('Sync-providers state row is missing');

const current = parseIsoTimestamp(row.lastAlertAt);
const requested = new Date(iso);
const latestIso = current && current > requested ? current.toISOString() : iso;
await tx
.update(ai_gateway_sync_providers_state)
.set({ stale_alert_last_posted_at: latestIso })
.where(eq(ai_gateway_sync_providers_state.id, 1));
return redisClient.set(SYNC_PROVIDERS_STALE_ALERT_LAST_POSTED_AT_REDIS_KEY, latestIso, {
ex: SYNC_PROVIDERS_STALE_ALERT_TTL_SECONDS,
});
return result;
});
}

export async function postStaleSyncAlert(input: {
Expand Down
33 changes: 23 additions & 10 deletions apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,18 @@ import { OpenRouterProvidersResponse } from '@/lib/ai-gateway/providers/openrout
import { fetchModelsForProvider } from '@/lib/ai-gateway/providers/openrouter/fetch-provider-models';
import { ai_gateway_sync_providers_state, modelsByProvider } from '@kilocode/db/schema';
import { db } from '@/lib/drizzle';
import { desc, lt, sql } from 'drizzle-orm';
import { desc, eq, lt, sql } from 'drizzle-orm';
import { captureException } from '@sentry/nextjs';
import { OPENROUTER, VERCEL_AI_GATEWAY } from '@/lib/ai-gateway/providers/provider-definitions';
import { logAutoModelChangesForAllOrgs } from '@/lib/organizations/auto-model-change-log';
import type { Provider } from '@/lib/ai-gateway/providers/types';
import type { StoredModel } from '@kilocode/db/schema-types';
import { EndpointsSchema, ModelsSchema } from '@kilocode/db/schema-types';
import { redisClient } from '@/lib/redis';
import { SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY } from '@/lib/redis-keys';
import {
AI_GATEWAY_STATE_REDIS_TTL_SECONDS,
SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY,
} from '@/lib/redis-keys';
import { syncDirectByokModels } from '@/lib/ai-gateway/providers/direct-byok/sync-direct-byok';
import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers';
import {
Expand Down Expand Up @@ -394,15 +397,25 @@ export async function syncAndStoreProviders() {
const direct_byok_model_counts = await syncDirectByokModels();
console.log('[syncAndStoreProviders] direct-byok model counts:', direct_byok_model_counts);

const completed_at = new Date().toISOString();
await redisClient.set(SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY, completed_at);
await db
.insert(ai_gateway_sync_providers_state)
.values({ last_completed_at: completed_at })
.onConflictDoUpdate({
target: ai_gateway_sync_providers_state.id,
set: { last_completed_at: completed_at },
const completed_at = await db.transaction(async tx => {
await tx.insert(ai_gateway_sync_providers_state).values({ id: 1 }).onConflictDoNothing();
const [row] = await tx
.select({ id: ai_gateway_sync_providers_state.id })
.from(ai_gateway_sync_providers_state)
.where(eq(ai_gateway_sync_providers_state.id, 1))
.for('update');
if (!row) throw new Error('Sync-providers state row is missing');

const completedAt = new Date().toISOString();
await tx
.update(ai_gateway_sync_providers_state)
.set({ last_completed_at: completedAt })
.where(eq(ai_gateway_sync_providers_state.id, 1));
await redisClient.set(SYNC_PROVIDERS_LAST_COMPLETED_AT_REDIS_KEY, completedAt, {
ex: AI_GATEWAY_STATE_REDIS_TTL_SECONDS,
});
return completedAt;
});

return {
id: result.id,
Expand Down
15 changes: 10 additions & 5 deletions apps/web/src/lib/ai-gateway/providers/routing-config.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@ import {
DEFAULT_VERCEL_PERCENTAGE_FREE,
GatewayRoutingConfigSchema,
} from '@/lib/ai-gateway/gateway-config';
import { redisClient } from '@/lib/redis';
import { VERCEL_ROUTING_REDIS_KEY } from '@/lib/redis-keys';
import { db } from '@/lib/drizzle';
import { ai_gateway_config } from '@kilocode/db/schema';
import { eq } from 'drizzle-orm';

export type RuntimeGatewayRoutingConfig = {
vercelPaid: number;
Expand All @@ -24,10 +25,14 @@ const DEFAULT_RUNTIME_GATEWAY_ROUTING_CONFIG: RuntimeGatewayRoutingConfig = {

export const getRuntimeGatewayRoutingConfig = createCachedFetch<RuntimeGatewayRoutingConfig>(
async () => {
const raw = await redisClient.get<string>(VERCEL_ROUTING_REDIS_KEY);
if (!raw) return DEFAULT_RUNTIME_GATEWAY_ROUTING_CONFIG;
const [row] = await db
.select({ config: ai_gateway_config.config })
.from(ai_gateway_config)
.where(eq(ai_gateway_config.id, 1))
.limit(1);
if (!row) return DEFAULT_RUNTIME_GATEWAY_ROUTING_CONFIG;
Comment thread
chrarnoldus marked this conversation as resolved.

const config = GatewayRoutingConfigSchema.parse(JSON.parse(raw));
const config = GatewayRoutingConfigSchema.parse(row.config);
return {
vercelPaid: config.vercel_routing_percentage ?? DEFAULT_VERCEL_PERCENTAGE,
vercelFree: config.vercel_routing_percentage_free ?? DEFAULT_VERCEL_PERCENTAGE_FREE,
Expand Down
136 changes: 65 additions & 71 deletions apps/web/src/lib/ai-gateway/request-logging-opt-ins.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@ import { ai_gateway_request_logging_opt_ins } from '@kilocode/db/schema';
import { createCachedFetch } from '@/lib/cached-fetch';
import { db } from '@/lib/drizzle';
import { redisClient } from '@/lib/redis';
import { REQUEST_LOGGING_OPT_INS_REDIS_KEY } from '@/lib/redis-keys';
import {
AI_GATEWAY_STATE_REDIS_TTL_SECONDS,
REQUEST_LOGGING_OPT_INS_REDIS_KEY,
} from '@/lib/redis-keys';
import { eq } from 'drizzle-orm';

export const RequestLoggingOptInSchema = z.object({
id: z.string().uuid(),
Expand All @@ -18,45 +22,6 @@ export const RequestLoggingOptInsSchema = z.array(RequestLoggingOptInSchema).max

export type RequestLoggingOptIn = z.infer<typeof RequestLoggingOptInSchema>;

const CREATE_OPT_IN_SCRIPT = `
local entries = {}
local raw = redis.call('GET', KEYS[1])
if raw then entries = cjson.decode(raw) end
local new_entry = cjson.decode(ARGV[1])
for _, entry in ipairs(entries) do
if entry.target_type == new_entry.target_type and entry.target_id == new_entry.target_id then
return 0
end
end
if #entries >= 500 then return -1 end
table.insert(entries, new_entry)
redis.call('SET', KEYS[1], cjson.encode(entries))
return 1
`;

const DELETE_OPT_IN_SCRIPT = `
local raw = redis.call('GET', KEYS[1])
if not raw then return 0 end
local entries = cjson.decode(raw)
local remaining = {}
local deleted = 0
for _, entry in ipairs(entries) do
if entry.id == ARGV[1] then
deleted = 1
else
table.insert(remaining, entry)
end
end
if deleted == 1 then
if #remaining == 0 then
redis.call('DEL', KEYS[1])
else
redis.call('SET', KEYS[1], cjson.encode(remaining))
end
end
return deleted
`;

const REQUEST_LOGGING_OPT_INS_CACHE_TTL_MS = process.env.NODE_ENV === 'test' ? 0 : 10_000;

export function hasMatchingRequestLoggingOptIn(
Expand All @@ -71,9 +36,12 @@ export function hasMatchingRequestLoggingOptIn(
}

export async function getRequestLoggingOptIns(): Promise<RequestLoggingOptIn[]> {
const raw = await redisClient.get<string>(REQUEST_LOGGING_OPT_INS_REDIS_KEY);
if (!raw) return [];
return RequestLoggingOptInsSchema.parse(JSON.parse(raw));
const [row] = await db
.select({ optIns: ai_gateway_request_logging_opt_ins.opt_ins })
.from(ai_gateway_request_logging_opt_ins)
.where(eq(ai_gateway_request_logging_opt_ins.id, 1))
.limit(1);
return RequestLoggingOptInsSchema.parse(row?.optIns ?? []);
Comment thread
chrarnoldus marked this conversation as resolved.
}

const getCachedRequestLoggingOptIns = createCachedFetch<RequestLoggingOptIn[]>(
Expand All @@ -82,43 +50,69 @@ const getCachedRequestLoggingOptIns = createCachedFetch<RequestLoggingOptIn[]>(
[]
);

async function mirrorRequestLoggingOptInsToDatabase(): Promise<void> {
const optIns = await getRequestLoggingOptIns();
await db
.insert(ai_gateway_request_logging_opt_ins)
.values({ opt_ins: optIns })
.onConflictDoUpdate({
target: ai_gateway_request_logging_opt_ins.id,
set: { opt_ins: optIns },
});
async function mirrorRequestLoggingOptInsToRedis(optIns: RequestLoggingOptIn[]): Promise<void> {
await redisClient.set(REQUEST_LOGGING_OPT_INS_REDIS_KEY, JSON.stringify(optIns), {
ex: AI_GATEWAY_STATE_REDIS_TTL_SECONDS,
});
}

export async function createRequestLoggingOptIn(
entry: RequestLoggingOptIn
): Promise<'created' | 'duplicate' | 'full'> {
const validated = RequestLoggingOptInSchema.parse(entry);
const result = await redisClient.eval<[string], number>(
CREATE_OPT_IN_SCRIPT,
[REQUEST_LOGGING_OPT_INS_REDIS_KEY],
[JSON.stringify(validated)]
);
if (result === 1) {
await mirrorRequestLoggingOptInsToDatabase();
return 'created';
}
if (result === 0) return 'duplicate';
return 'full';
return db.transaction(async tx => {
await tx
.insert(ai_gateway_request_logging_opt_ins)
.values({ opt_ins: [] })
.onConflictDoNothing();
const [row] = await tx
.select({ optIns: ai_gateway_request_logging_opt_ins.opt_ins })
.from(ai_gateway_request_logging_opt_ins)
.where(eq(ai_gateway_request_logging_opt_ins.id, 1))
.for('update');
if (!row) throw new Error('Request logging opt-in state row is missing');

const optIns = RequestLoggingOptInsSchema.parse(row.optIns);
if (
optIns.some(
optIn =>
optIn.target_type === validated.target_type && optIn.target_id === validated.target_id
)
) {
return 'duplicate' as const;
}
if (optIns.length >= 500) return 'full' as const;

const updatedOptIns = [...optIns, validated];
await tx
.update(ai_gateway_request_logging_opt_ins)
.set({ opt_ins: updatedOptIns })
.where(eq(ai_gateway_request_logging_opt_ins.id, 1));
await mirrorRequestLoggingOptInsToRedis(updatedOptIns);
return 'created' as const;
});
}

export async function deleteRequestLoggingOptIn(id: string): Promise<boolean> {
const result = await redisClient.eval<[string], number>(
DELETE_OPT_IN_SCRIPT,
[REQUEST_LOGGING_OPT_INS_REDIS_KEY],
[id]
);
if (result !== 1) return false;
await mirrorRequestLoggingOptInsToDatabase();
return true;
return db.transaction(async tx => {
const [row] = await tx
.select({ optIns: ai_gateway_request_logging_opt_ins.opt_ins })
.from(ai_gateway_request_logging_opt_ins)
.where(eq(ai_gateway_request_logging_opt_ins.id, 1))
.for('update');
if (!row) return false;

const optIns = RequestLoggingOptInsSchema.parse(row.optIns);
const remaining = optIns.filter(entry => entry.id !== id);
if (remaining.length === optIns.length) return false;

await tx
.update(ai_gateway_request_logging_opt_ins)
.set({ opt_ins: remaining })
.where(eq(ai_gateway_request_logging_opt_ins.id, 1));
await mirrorRequestLoggingOptInsToRedis(remaining);
return true;
});
}

export async function isDynamicallyOptedIntoRequestLogging(params: {
Expand Down
Loading