From a70001f461b1065103be7f877fdbcb8ec9ef7e91 Mon Sep 17 00:00:00 2001 From: RainbowBird Date: Wed, 4 Mar 2026 22:55:07 +0800 Subject: [PATCH] fix(server): resolve config kv properly --- apps/server/src/middlewares/config-guard.ts | 21 ++++++ apps/server/src/routes/stripe.ts | 5 +- apps/server/src/routes/v1completions.ts | 3 +- .../src/services/__test__/config-kv.test.ts | 74 +++++++++++++++++++ .../server/src/services/__test__/flux.test.ts | 15 +++- apps/server/src/services/config-kv.ts | 12 ++- 6 files changed, 124 insertions(+), 6 deletions(-) create mode 100644 apps/server/src/middlewares/config-guard.ts create mode 100644 apps/server/src/services/__test__/config-kv.test.ts diff --git a/apps/server/src/middlewares/config-guard.ts b/apps/server/src/middlewares/config-guard.ts new file mode 100644 index 000000000..003282461 --- /dev/null +++ b/apps/server/src/middlewares/config-guard.ts @@ -0,0 +1,21 @@ +import type { MiddlewareHandler } from 'hono' + +import type { ConfigKVService } from '../services/config-kv' +import type { HonoEnv } from '../types/hono' + +import { createServiceUnavailableError } from '../utils/error' + +/** + * Middleware factory that checks required config keys exist in Redis. + * Returns 503 if any key is missing. + */ +export function configGuard(configKV: ConfigKVService, keys: Parameters[0][]): MiddlewareHandler { + return async (_c, next) => { + for (const key of keys) { + const value = await configKV.getOptional(key) + if (value === null) + throw createServiceUnavailableError(`Config key "${key}" is not set in Redis`, 'CONFIG_NOT_SET') + } + await next() + } +} diff --git a/apps/server/src/routes/stripe.ts b/apps/server/src/routes/stripe.ts index a4f2d907c..91389b9ef 100644 --- a/apps/server/src/routes/stripe.ts +++ b/apps/server/src/routes/stripe.ts @@ -10,6 +10,7 @@ import { Hono } from 'hono' import { integer, minValue, number, object, pipe, safeParse } from 'valibot' import { authGuard } from '../middlewares/auth' +import { configGuard } from '../middlewares/config-guard' import { createBadRequestError, createServiceUnavailableError } from '../utils/error' const CheckoutBodySchema = object({ @@ -19,8 +20,10 @@ const CheckoutBodySchema = object({ export function createStripeRoutes(fluxService: FluxService, stripeService: StripeService, configKV: ConfigKVService, env: Env) { const stripe = env.STRIPE_SECRET_KEY ? new Stripe(env.STRIPE_SECRET_KEY) : null + const fluxConfigGuard = configGuard(configKV, ['FLUX_PER_CENT']) + return new Hono() - .post('/checkout', authGuard, async (c) => { + .post('/checkout', authGuard, fluxConfigGuard, async (c) => { if (!stripe) throw createServiceUnavailableError('Stripe is not configured', 'STRIPE_NOT_CONFIGURED') diff --git a/apps/server/src/routes/v1completions.ts b/apps/server/src/routes/v1completions.ts index 82b05612f..463b8091f 100644 --- a/apps/server/src/routes/v1completions.ts +++ b/apps/server/src/routes/v1completions.ts @@ -8,6 +8,7 @@ import type { HonoEnv } from '../types/hono' import { Hono } from 'hono' import { authGuard } from '../middlewares/auth' +import { configGuard } from '../middlewares/config-guard' import { createPaymentRequiredError } from '../utils/error' // Only forward these headers from the upstream LLM response @@ -53,7 +54,7 @@ export function createV1CompletionsRoutes(fluxService: FluxService, configKV: Co } return new Hono() - .use('*', authGuard) + .use('*', authGuard, configGuard(configKV, ['FLUX_PER_REQUEST'])) .post('/chat/completions', handleCompletion) .post('/chat/completion', handleCompletion) } diff --git a/apps/server/src/services/__test__/config-kv.test.ts b/apps/server/src/services/__test__/config-kv.test.ts new file mode 100644 index 000000000..1d75cf9cc --- /dev/null +++ b/apps/server/src/services/__test__/config-kv.test.ts @@ -0,0 +1,74 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { createConfigKVService } from '../config-kv' + +function createMockRedis() { + const store = new Map() + return { + get: vi.fn(async (key: string) => store.get(key) ?? null), + set: vi.fn(async (key: string, value: string) => { store.set(key, value) }), + _store: store, + } +} + +describe('configKVService', () => { + let redis: ReturnType + let service: ReturnType + + beforeEach(() => { + redis = createMockRedis() + service = createConfigKVService(redis as any) + }) + + // --- get --- + + it('get should throw 503 when key is not set', async () => { + await expect(service.get('FLUX_PER_CENT')) + .rejects + .toThrow('Config key "FLUX_PER_CENT" is not set in Redis') + }) + + it('get should return numeric value when key is set', async () => { + redis._store.set('config:FLUX_PER_CENT', '5') + + const value = await service.get('FLUX_PER_CENT') + expect(value).toBe(5) + }) + + it('get should read from correct prefixed key', async () => { + redis._store.set('config:FLUX_PER_REQUEST', '3') + + await service.get('FLUX_PER_REQUEST') + expect(redis.get).toHaveBeenCalledWith('config:FLUX_PER_REQUEST') + }) + + // --- getOptional --- + + it('getOptional should return null when key is not set', async () => { + const value = await service.getOptional('FLUX_PER_CENT') + expect(value).toBeNull() + }) + + it('getOptional should return numeric value when key is set', async () => { + redis._store.set('config:INITIAL_USER_FLUX', '200') + + const value = await service.getOptional('INITIAL_USER_FLUX') + expect(value).toBe(200) + }) + + // --- set --- + + it('set should write value to Redis with prefix', async () => { + await service.set('FLUX_PER_CENT', 10) + + expect(redis.set).toHaveBeenCalledWith('config:FLUX_PER_CENT', '10') + expect(redis._store.get('config:FLUX_PER_CENT')).toBe('10') + }) + + it('set then get should round-trip correctly', async () => { + await service.set('INITIAL_USER_FLUX', 500) + + const value = await service.get('INITIAL_USER_FLUX') + expect(value).toBe(500) + }) +}) diff --git a/apps/server/src/services/__test__/flux.test.ts b/apps/server/src/services/__test__/flux.test.ts index 53169d02f..7e62ec521 100644 --- a/apps/server/src/services/__test__/flux.test.ts +++ b/apps/server/src/services/__test__/flux.test.ts @@ -1,10 +1,21 @@ -import { beforeAll, describe, expect, it } from 'vitest' +import type { createConfigKVService } from '../config-kv' + +import { beforeAll, describe, expect, it, vi } from 'vitest' import { mockDB } from '../../libs/mock-db' import { createFluxService } from '../flux' import * as schema from '../../schemas' +function createMockConfigKV(overrides: Record = {}): ReturnType { + const defaults: Record = { INITIAL_USER_FLUX: 100, FLUX_PER_CENT: 1, FLUX_PER_REQUEST: 1, ...overrides } + return { + get: vi.fn(async (key: string) => defaults[key]), + getOptional: vi.fn(async (key: string) => defaults[key] ?? null), + set: vi.fn(), + } as any +} + describe('fluxService', () => { let db: any let service: ReturnType @@ -12,7 +23,7 @@ describe('fluxService', () => { beforeAll(async () => { db = await mockDB(schema) - service = createFluxService(db) + service = createFluxService(db, createMockConfigKV()) // Create a test user for foreign key constraints const [user] = await db.insert(schema.user).values({ diff --git a/apps/server/src/services/config-kv.ts b/apps/server/src/services/config-kv.ts index 6a00ea98c..d25c9fa35 100644 --- a/apps/server/src/services/config-kv.ts +++ b/apps/server/src/services/config-kv.ts @@ -12,14 +12,22 @@ const KEY_PREFIX = 'config:' export function createConfigKVService(redis: Redis) { return { - async get(key: K): Promise { + async getOptional(key: K): Promise { const raw = await redis.get(`${KEY_PREFIX}${key}`) if (raw === null) - throw createServiceUnavailableError(`Config key "${key}" is not set in Redis`, 'CONFIG_NOT_SET') + return null return Number(raw) as ConfigDefinitions[K] }, + async get(key: K): Promise { + const value = await this.getOptional(key) + if (value === null) + throw createServiceUnavailableError(`Config key "${key}" is not set in Redis`, 'CONFIG_NOT_SET') + + return value + }, + async set(key: K, value: ConfigDefinitions[K]): Promise { await redis.set(`${KEY_PREFIX}${key}`, String(value)) },