From 14fb1837231b8481da65a2eaa167634807c08071 Mon Sep 17 00:00:00 2001 From: RainbowBird Date: Wed, 1 Jul 2026 22:13:28 +0800 Subject: [PATCH] fix(server): gate streaming tts requests --- .../src/routes/audio-speech-ws/index.ts | 7 +- .../src/routes/audio-speech-ws/route.test.ts | 96 +++++++++++- .../src/routes/audio-speech-ws/session.ts | 138 +++++++++++++++++- 3 files changed, 227 insertions(+), 14 deletions(-) diff --git a/apps/server/src/routes/audio-speech-ws/index.ts b/apps/server/src/routes/audio-speech-ws/index.ts index 79bd18839..42d6cee75 100644 --- a/apps/server/src/routes/audio-speech-ws/index.ts +++ b/apps/server/src/routes/audio-speech-ws/index.ts @@ -23,8 +23,8 @@ export type { AudioSpeechWsHandlersOptions } from './types' * Expects: * - The route handler has already resolved auth via the `?token=` query * (see app.ts wiring) and passes a verified `userId` in. - * - `UNSPEECH_UPSTREAM.streaming` configKV subtree is populated with at least - * one key; absent config rejects the upgrade with policy-violation close. + * - The client sends a `start` control frame first. The session validates the + * requested streaming model and voice before dialing upstream. * * Returns: * - A function that takes `userId` and returns hono `WSEvents`. Each call @@ -38,9 +38,6 @@ export function createAudioSpeechWsHandlers(opts: AudioSpeechWsHandlersOptions) return { onOpen(_event, ws) { sessionState.attachClient(ws) - // Dial upstream inside the open handler so failure surfaces as a - // clean close on the client ws rather than a 500 on the upgrade. - void sessionState.dialUpstream() }, onMessage(message, ws) { sessionState.handleClientMessage(message, ws) diff --git a/apps/server/src/routes/audio-speech-ws/route.test.ts b/apps/server/src/routes/audio-speech-ws/route.test.ts index 9345e9f2b..9523fa6a2 100644 --- a/apps/server/src/routes/audio-speech-ws/route.test.ts +++ b/apps/server/src/routes/audio-speech-ws/route.test.ts @@ -12,6 +12,7 @@ import { createAudioSpeechWsHandlers } from './index' interface MockUpstream { url: string + restBaseURL: string /** Outgoing JSON frames the server should send after receiving `start`. */ scriptedResponses: Array< | { kind: 'json', payload: Record } @@ -24,11 +25,22 @@ interface MockUpstream { close: () => Promise } -async function startMockUpstream(scriptedResponses: MockUpstream['scriptedResponses']): Promise { +async function startMockUpstream( + scriptedResponses: MockUpstream['scriptedResponses'], + voices: Array<{ id: string, name?: string }> = [{ id: 'mock', name: 'Mock Voice' }], +): Promise { const receivedFrames: MockUpstream['receivedFrames'] = [] let observedAuth: string | undefined - const httpServer = createServer() + const httpServer = createServer((req, res) => { + if (req.url?.startsWith('/api/voices')) { + res.writeHead(200, { 'content-type': 'application/json' }) + res.end(JSON.stringify({ voices })) + return + } + res.writeHead(404, { 'content-type': 'application/json' }) + res.end(JSON.stringify({ error: 'not_found' })) + }) const wss = new WebSocketServer({ server: httpServer }) wss.on('connection', (ws, req) => { @@ -91,6 +103,7 @@ async function startMockUpstream(scriptedResponses: MockUpstream['scriptedRespon return { url: `ws://127.0.0.1:${port}`, + restBaseURL: `http://127.0.0.1:${port}`, scriptedResponses, receivedFrames, get observedAuth() { @@ -148,8 +161,10 @@ function makeMockClientWs(): MockClientWs { function makeFakeDeps(overrides: { upstreamURL: string + restBaseURL?: string fluxBalance: number decryptedKey?: string + streamingModels?: Array<{ id: string, name?: string, description?: string }> }) { const ttsMeter = { assertCanAfford: vi.fn(async (_userId: string, _newUnits: number, currentBalance: number) => { @@ -177,11 +192,15 @@ function makeFakeDeps(overrides: { getOptional: vi.fn(async (key: string) => { if (key === 'UNSPEECH_UPSTREAM') { return { - restBaseURL: 'http://unspeech.local:5933', + restBaseURL: overrides.restBaseURL ?? 'http://unspeech.local:5933', streaming: { baseURL: overrides.upstreamURL, keys: [{ id: 'test-key-1', ciphertext: 'ENCRYPTED_PLACEHOLDER' }], adapterParams: {}, + models: overrides.streamingModels ?? [ + { id: 'volcengine/seed-tts-1.0', name: 'Seed-TTS 1.0' }, + { id: 'volcengine/seed-tts-2.0', name: 'Seed-TTS 2.0' }, + ], }, } } @@ -228,7 +247,7 @@ describe('audio-speech-ws route', () => { { kind: 'json', payload: { event: 'session.finished', payload: { usage: { text_words: 42 } } } }, ]) - const deps = makeFakeDeps({ upstreamURL: upstream.url, fluxBalance: 100 }) + const deps = makeFakeDeps({ upstreamURL: upstream.url, restBaseURL: upstream.restBaseURL, fluxBalance: 100 }) const handlers = createAudioSpeechWsHandlers(deps as any) const events = handlers('user-123', { voiceType: 'official_selected' }) const client = makeMockClientWs() @@ -294,12 +313,14 @@ describe('audio-speech-ws route', () => { it('refuses the session with insufficient_flux when the user is broke', async () => { upstream = await startMockUpstream([]) - const deps = makeFakeDeps({ upstreamURL: upstream.url, fluxBalance: 0 }) + const deps = makeFakeDeps({ upstreamURL: upstream.url, restBaseURL: upstream.restBaseURL, fluxBalance: 0 }) const handlers = createAudioSpeechWsHandlers(deps as any) const events = handlers('user-broke', { trigger: 'auto', source: 'chat_auto_tts' }) const client = makeMockClientWs() - await driveClientSession(events, client, []) + await driveClientSession(events, client, [ + JSON.stringify({ event: 'start', model: 'volcengine/seed-tts-2.0', voice: 'mock' }), + ]) // Upstream should never have been dialed — pre-flight fails first. expect(upstream.receivedFrames).toHaveLength(0) @@ -337,7 +358,9 @@ describe('audio-speech-ws route', () => { const events = handlers('user-noconf') const client = makeMockClientWs() - await driveClientSession(events, client, []) + await driveClientSession(events, client, [ + JSON.stringify({ event: 'start', model: 'volcengine/seed-tts-2.0', voice: 'mock' }), + ]) const errorFrame = client.sent.find(s => s.kind === 'text') expect(errorFrame).toBeDefined() @@ -348,6 +371,63 @@ describe('audio-speech-ws route', () => { expect(client.closed).toBe(true) }) + it('refuses an unconfigured streaming model before dialing upstream', async () => { + upstream = await startMockUpstream([]) + const deps = makeFakeDeps({ + upstreamURL: upstream.url, + restBaseURL: upstream.restBaseURL, + fluxBalance: 100, + streamingModels: [{ id: 'volcengine/seed-tts-2.0', name: 'Seed-TTS 2.0' }], + }) + const handlers = createAudioSpeechWsHandlers(deps as any) + const events = handlers('user-disabled-model') + const client = makeMockClientWs() + + await driveClientSession(events, client, [ + JSON.stringify({ event: 'start', model: 'volcengine/seed-tts-disabled', voice: 'mock' }), + JSON.stringify({ event: 'text', text: 'must not leak upstream' }), + JSON.stringify({ event: 'finish' }), + ]) + await new Promise(r => setTimeout(r, 100)) + + expect(upstream.observedAuth).toBeUndefined() + expect(upstream.receivedFrames).toHaveLength(0) + const errorFrame = client.sent.find(s => s.kind === 'text') + expect(errorFrame).toBeDefined() + expect(JSON.parse(errorFrame!.data as string)).toMatchObject({ + event: 'error', + code: 'streaming_tts_model_not_enabled', + }) + expect(client.closed).toBe(true) + expect(client.closeCode).toBe(1008) + }) + + it('refuses an unknown streaming voice before dialing upstream', async () => { + upstream = await startMockUpstream([], [{ id: 'enabled-voice', name: 'Enabled Voice' }]) + const deps = makeFakeDeps({ upstreamURL: upstream.url, restBaseURL: upstream.restBaseURL, fluxBalance: 100 }) + const handlers = createAudioSpeechWsHandlers(deps as any) + const events = handlers('user-disabled-voice') + const client = makeMockClientWs() + + await driveClientSession(events, client, [ + JSON.stringify({ event: 'start', model: 'volcengine/seed-tts-2.0', voice: 'disabled-voice' }), + JSON.stringify({ event: 'text', text: 'must not leak upstream' }), + JSON.stringify({ event: 'finish' }), + ]) + await new Promise(r => setTimeout(r, 100)) + + expect(upstream.observedAuth).toBeUndefined() + expect(upstream.receivedFrames).toHaveLength(0) + const errorFrame = client.sent.find(s => s.kind === 'text') + expect(errorFrame).toBeDefined() + expect(JSON.parse(errorFrame!.data as string)).toMatchObject({ + event: 'error', + code: 'streaming_tts_voice_not_enabled', + }) + expect(client.closed).toBe(true) + expect(client.closeCode).toBe(1008) + }) + it('falls back to input-char count for billing when upstream omits usage', async () => { // No usage in session.finished — proxy must bill the cumulative // length of every `text` frame's `text` field instead. @@ -357,7 +437,7 @@ describe('audio-speech-ws route', () => { { kind: 'json', payload: { event: 'session.finished', payload: {} } }, ]) - const deps = makeFakeDeps({ upstreamURL: upstream.url, fluxBalance: 100 }) + const deps = makeFakeDeps({ upstreamURL: upstream.url, restBaseURL: upstream.restBaseURL, fluxBalance: 100 }) const handlers = createAudioSpeechWsHandlers(deps as any) const events = handlers('user-no-usage') const client = makeMockClientWs() diff --git a/apps/server/src/routes/audio-speech-ws/session.ts b/apps/server/src/routes/audio-speech-ws/session.ts index f62ed2a17..52542c51e 100644 --- a/apps/server/src/routes/audio-speech-ws/session.ts +++ b/apps/server/src/routes/audio-speech-ws/session.ts @@ -10,6 +10,7 @@ import WebSocket from 'ws' import { useLogger } from '@guiiai/logg' import { context as otelContext, SpanStatusCode, trace } from '@opentelemetry/api' +import { ofetch } from 'ofetch' import { fluxBalanceBucket } from '../../services/domain/flux-balance' import { ApiError } from '../../utils/error' @@ -43,7 +44,7 @@ const tracer = trace.getTracer('audio-speech-ws') export interface AudioSpeechSessionState { /** Stores the accepted client websocket. */ attachClient: (ws: WSContext) => void - /** Reads config, checks balance, decrypts the upstream key, and dials upstream. */ + /** Reads config, checks balance, decrypts the upstream key, and dials upstream after the start frame is accepted. */ dialUpstream: () => Promise /** Forwards a client frame or queues it while the upstream connection opens. */ handleClientMessage: (message: { data: unknown }, ws: WSContext) => void @@ -94,6 +95,9 @@ export function createSessionState( let upstreamReady = false let closed = false let billed = false + let startFrameAccepted = false + let startValidationStarted = false + let dialStarted = false let totalInputChars = 0 let preflightFluxBalance: number | undefined let modelLabel = STREAM_MODEL_LABEL_FALLBACK @@ -110,6 +114,10 @@ export function createSessionState( } async function dialUpstream() { + if (dialStarted) + return + dialStarted = true + void opts.productEventService.track({ userId, feature: 'tts', @@ -262,6 +270,31 @@ export function createSessionState( ? Buffer.from(message.data) : Buffer.from(message.data as ArrayBufferLike) + if (!startValidationStarted) { + if (isBinary || typeof payload !== 'string') { + closeWithError(1008, 'invalid_start_frame') + return + } + + const startFrame = parseStartFrame(payload) + if (!startFrame) { + closeWithError(1008, 'invalid_start_frame') + return + } + + startValidationStarted = true + modelLabel = startFrame.model + voiceLabel = startFrame.voice + pendingClientFrames.push({ data: payload, isBinary }) + void validateStartFrame(startFrame).then((accepted) => { + if (!accepted || closed) + return + startFrameAccepted = true + void dialUpstream() + }) + return + } + // Sniff input chars from text frames so billing has a fallback when // upstream usage.text_words is absent. Only the `text` event contributes; // start/finish/cancel do not. @@ -269,6 +302,11 @@ export function createSessionState( maybeAccountInputChars(payload) } + if (!startFrameAccepted && !dialStarted) { + pendingClientFrames.push({ data: payload, isBinary }) + return + } + if (!upstreamWs || !upstreamReady) { pendingClientFrames.push({ data: payload, isBinary }) return @@ -380,6 +418,55 @@ export function createSessionState( } } + async function validateStartFrame(frame: StreamingTtsStartFrame): Promise { + let unspeech: Awaited> + try { + unspeech = await opts.configKV.getOptional('UNSPEECH_UPSTREAM') + } + catch (err) { + log.withError(err).error('UNSPEECH_UPSTREAM read failed before streaming tts start') + closeWithError(1011, 'config_unavailable') + return false + } + + const upstreamConfig = unspeech?.streaming + if (!unspeech?.restBaseURL || !upstreamConfig?.baseURL || upstreamConfig.keys.length === 0) { + closeWithError(1008, 'streaming_tts_not_configured') + return false + } + + const configuredModels = upstreamConfig.models ?? [] + if (!configuredModels.some((model: { id: string }) => model.id === frame.model)) { + closeWithError(1008, 'streaming_tts_model_not_enabled') + return false + } + + const resourceId = streamingModelResourceId(frame.model) + const voicesURL = streamingVoicesURL(unspeech.restBaseURL, resourceId) + if (!voicesURL) { + closeWithError(1011, 'streaming_tts_voice_catalog_unavailable') + return false + } + + let data: { voices?: unknown[] } + try { + data = await ofetch(voicesURL, { timeout: 5000 }) as { voices?: unknown[] } + } + catch (err) { + log.withError(err).withFields({ voicesURL }).warn('streaming tts voice catalog fetch failed') + closeWithError(1011, 'streaming_tts_voice_catalog_unavailable') + return false + } + + const voices = Array.isArray(data.voices) ? data.voices : [] + if (!voices.some(voice => streamingVoiceId(voice) === frame.voice)) { + closeWithError(1008, 'streaming_tts_voice_not_enabled') + return false + } + + return true + } + async function billSession(units: number, reason: string) { if (billed) return @@ -600,3 +687,52 @@ function isPaymentRequiredError(err: unknown): boolean { && 'statusCode' in err && (err as { statusCode?: unknown }).statusCode === 402 } + +interface StreamingTtsStartFrame { + event: 'start' + model: string + voice: string +} + +function parseStartFrame(rawText: string): StreamingTtsStartFrame | null { + try { + const parsed = JSON.parse(rawText) as Record + if (parsed.event !== 'start') + return null + if (typeof parsed.model !== 'string' || parsed.model.length === 0) + return null + if (typeof parsed.voice !== 'string' || parsed.voice.length === 0) + return null + return { + event: 'start', + model: parsed.model, + voice: parsed.voice, + } + } + catch { + return null + } +} + +function streamingModelResourceId(model: string): string { + return model.includes('/') ? model.split('/', 2)[1] : model +} + +function streamingVoicesURL(restBaseURL: string, resourceId: string): string | null { + try { + const url = new URL(restBaseURL) + url.pathname = '/api/voices' + url.search = new URLSearchParams({ provider: 'volcengine', model: resourceId }).toString() + return url.toString() + } + catch { + return null + } +} + +function streamingVoiceId(voice: unknown): string | null { + if (typeof voice !== 'object' || voice == null) + return null + const id = (voice as { id?: unknown }).id + return typeof id === 'string' && id.length > 0 ? id : null +}