diff --git a/apps/web/src/lib/agent-harness/model-stream.test.ts b/apps/web/src/lib/agent-harness/model-stream.test.ts new file mode 100644 index 0000000000..ff7e990c65 --- /dev/null +++ b/apps/web/src/lib/agent-harness/model-stream.test.ts @@ -0,0 +1,230 @@ +import { expect, it } from '@jest/globals'; +import { z } from 'zod'; +import { TRPCError } from '@trpc/server'; +import { getHTTPStatusCodeFromError } from '@trpc/server/http'; +import { createOpenAICompatible } from '@ai-sdk/openai-compatible'; +import { streamText } from 'ai'; +import { boundedBody, streamHarnessModel } from './model-stream'; + +const outputHeaders = { + 'content-type': 'Text/Event-Stream; charset=utf-8', + 'request-id': 'request-1', + authorization: 'upstream-secret', +}; +const sse = (body: BodyInit | null) => new Response(body, { headers: outputHeaders }); +function invoke( + upstream: Response, + request = new AbortController(), + deadline = new AbortController() +) { + const abort = new AbortController(); + const signal = AbortSignal.any([request.signal, deadline.signal, abort.signal]); + const response = streamHarnessModel( + upstream, + new Headers({ 'cache-control': 'no-store', 'content-type': 'application/json' }), + { + signal, + abort, + failure: code => ({ + error: { code, message: 'Safe error', retryable: code === 429 || code >= 500 }, + }), + errorStatus: (error, invalid) => { + if (request.signal.aborted) return 499; + if (error instanceof TRPCError) return getHTTPStatusCodeFromError(error); + return error instanceof SyntaxError || error instanceof z.ZodError ? invalid : 503; + }, + } + ); + return { response, abort }; +} + +it.each(['', 'x-provider: private-field\nretry: private-retry\n\n'])( + 'streams multiline SDK text, reasoning, usage, and identifiers (prefix %j)', + async prefix => { + const chunk = { + id: 'generation-1', + model: 'paid/model', + created: 1, + choices: [ + { + index: 0, + delta: { content: 'Hello', reasoning_content: 'Thinking' }, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5, cost: 0.002 }, + }; + let upstream!: ReadableStreamDefaultController; + const cancelled = Promise.withResolvers(); + const { response, abort } = invoke( + sse( + new ReadableStream({ + start(controller) { + upstream = controller; + const data = JSON.stringify(chunk, null, 2).replaceAll('\n', '\ndata: '); + controller.enqueue(new TextEncoder().encode(`${prefix}data: ${data}\n\n`)); + }, + cancel: () => cancelled.resolve(), + }) + ) + ); + const reply = await response; + const raw = reply.clone().text(); + const provider = createOpenAICompatible({ + name: 'harness', + baseURL: 'https://unused.example', + fetch: async () => reply, + }); + const result = streamText({ + model: provider('paid/model'), + messages: [{ role: 'user', content: 'Hello' }], + maxOutputTokens: 128, + maxRetries: 0, + }); + const reader = result.textStream.getReader(); + expect((await reader.read()).value).toBe('Hello'); + upstream.enqueue(new TextEncoder().encode('data: [DONE]\n\n')); + expect((await reader.read()).done).toBe(true); + expect((await result.response).id).toBe('generation-1'); + expect(await result.reasoningText).toBe('Thinking'); + expect(await result.usage).toMatchObject({ inputTokens: 3, outputTokens: 2 }); + expect(await result.finishReason).toBe('stop'); + expect(await raw).toBe(`data: ${JSON.stringify(chunk)}\n\ndata: [DONE]\n\n`); + expect(reply.headers.get('authorization')).toBeNull(); + expect(reply.headers.get('request-id')).toBe('request-1'); + expect(reply.headers.get('cache-control')).toBe('no-store'); + expect(reply.headers.get('content-encoding')).toBe('identity'); + await cancelled.promise; + expect(abort.signal.aborted).toBe(true); + } +); +it('preserves split UTF-8, tool extensions, and separate usage events', async () => { + const tool = { + id: 'call-1', + function: { name: 'lookup', arguments: '{}' }, + extra_content: { google: { thought_signature: 'signature-1' } }, + }; + const delta = { + choices: [{ delta: { content: '界', reasoning: 'Thinking', tool_calls: [tool] } }], + }; + const usage = { + choices: [], + usage: { + prompt_tokens: 3, + completion_tokens: 2, + cost: 0.002, + prompt_tokens_details: { cached_tokens: 1 }, + completion_tokens_details: { reasoning_tokens: 1, accepted_prediction_tokens: 1 }, + }, + }; + const data = `data: ${JSON.stringify(delta)}\n\ndata: ${JSON.stringify(usage)}\n\ndata: [DONE]\n\n`; + const bytes = new TextEncoder().encode(data); + const split = bytes.indexOf(0xe7) + 1; + const body = new ReadableStream({ + start(controller) { + controller.enqueue(bytes.slice(0, split)); + controller.enqueue(bytes.slice(split, split + 1)); + controller.enqueue(bytes.slice(split + 1)); + controller.close(); + }, + }); + expect(await (await invoke(sse(body)).response).text()).toBe(data); +}); +it.each([400, 401, 402, 403, 429, 503])('sanitizes terminal stream error %s', async code => { + const error = JSON.stringify({ + error: { code, message: 'upstream-secret', metadata: { authorization: 'credential' } }, + }); + const { response, abort } = invoke(sse(`data: ${error}\n\ndata: [DONE]\n\n`)); + const reply = await response; + const text = await reply.text(); + expect(reply.status).toBe(200); + expect(text).toContain(`"code":${code}`); + expect(text).toContain(`"retryable":${code === 429 || code === 503}`); + expect(text).not.toMatch(/upstream-secret|credential|\[DONE\]/); + expect(abort.signal.aborted).toBe(true); +}); +it.each([ + ['data: private-malformed\n\n', 422], + [`data: ${'x'.repeat(1024 * 1024)}\n\n`, 413], + ['data: []\n\n', 422], + ['data: {"metadata":"private-malformed"}\n\n', 422], + ['event: error\ndata: {"choices":[]}\n\n', 422], + ['data: {"error":{"code":200,"message":"private-malformed"}}\n\n', 422], +] as const)('rejects malformed or oversized streams (case %#)', async (data, code) => { + const text = await (await invoke(sse(data)).response).text(); + expect(text).toContain(`"code":${code}`); + expect(text).not.toContain('private-malformed'); +}); +it.each([ + { choices: ['private-malformed'] }, + { choices: [{ delta: 'private-malformed' }] }, + { choices: [{ delta: { content: ['private-malformed'] } }] }, + { choices: [{ delta: { tool_calls: ['private-malformed'] } }] }, + { usage: 'private-malformed' }, + { usage: ['private-malformed'] }, + { usage: { prompt_tokens: 'private-malformed' } }, + { usage: { completion_tokens_details: { reasoning_tokens: 'private-malformed' } } }, +])('sanitizes malformed SDK payloads before forwarding (case %#)', async payload => { + const data = JSON.stringify({ choices: [], ...payload }); + const { response, abort } = invoke(sse(`data: ${data}\n\ndata: [DONE]\n\n`)); + const text = await (await response).text(); + expect(text).not.toContain('private-malformed'); + expect(text).toBe('data: {"error":{"code":422,"message":"Safe error","retryable":false}}\n\n'); + expect(abort.signal.aborted).toBe(true); +}); +it.each([0, 1])('counts cumulative bytes before decoding (extra bytes: %s)', async extra => { + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(512 * 1024)); + controller.enqueue(new Uint8Array(512 * 1024 + extra)); + controller.close(); + }, + }); + const result = new Response(boundedBody(body, new AbortController().signal)).arrayBuffer(); + if (extra) await expect(result).rejects.toMatchObject({ code: 'PAYLOAD_TOO_LARGE' }); + else expect((await result).byteLength).toBe(1024 * 1024); +}); +it.each(['', 'data: [DONE]\n\n'])('preserves empty output %j', async data => { + const { response, abort } = invoke(sse(data)); + expect(await (await response).text()).toBe(data); + expect(abort.signal.aborted).toBe(true); +}); +it.each([ + new Response('upstream-secret', { status: 429 }), + new Response('private', { headers: { 'content-type': 'text/html' } }), + sse(null), +])('rejects HTTP failures, wrong media, and missing bodies (case %#)', async upstream => { + const reply = await invoke(upstream).response; + expect(reply.status).toBe(upstream.status === 429 ? 429 : 422); + expect(await reply.text()).not.toMatch(/upstream-secret|private/); +}); +it('sanitizes a failed upstream read', async () => { + const body = new ReadableStream({ + start(controller) { + controller.error(new Error('transport-secret')); + }, + }); + const { response, abort } = invoke(sse(body)); + const text = await (await response).text(); + expect(text).toContain('"code":503'); + expect(text).not.toContain('transport-secret'); + expect(abort.signal.aborted).toBe(true); +}); +it.each(['request', 'reader', 'deadline'])('cancels and cleans up through the %s', async target => { + const cancelled = Promise.withResolvers(); + const body = new ReadableStream({ cancel: () => cancelled.resolve() }); + const request = new AbortController(); + const deadline = new AbortController(); + const { response, abort } = invoke(sse(body), request, deadline); + const reply = await response; + if (target === 'reader') await reply.body!.cancel(); + else { + const text = reply.text(); + (target === 'request' ? request : deadline).abort(new Error('private')); + const output = await text; + expect(output).toContain(`"code":${target === 'request' ? 499 : 503}`); + expect(output).not.toContain('private'); + } + await cancelled.promise; + expect(abort.signal.aborted).toBe(true); +}); diff --git a/apps/web/src/lib/agent-harness/model-stream.ts b/apps/web/src/lib/agent-harness/model-stream.ts new file mode 100644 index 0000000000..5e09575ac3 --- /dev/null +++ b/apps/web/src/lib/agent-harness/model-stream.ts @@ -0,0 +1,153 @@ +import 'server-only'; +import { z } from 'zod'; +import { TRPCError } from '@trpc/server'; +import { EventSourceParserStream } from 'eventsource-parser/stream'; + +export function boundedBody(body: ReadableStream, signal: AbortSignal) { + let bytes = 0; + return body.pipeThrough( + new TransformStream({ + transform(chunk, controller) { + bytes += chunk.byteLength; + if (bytes > 1024 * 1024) throw new TRPCError({ code: 'PAYLOAD_TOO_LARGE' }); + controller.enqueue(chunk); + }, + }), + { signal } + ); +} +const StreamEnvelope = z.object({ error: z.unknown().optional() }); +// Match @ai-sdk/openai-compatible's unexported chunk schema. Forward the original +// JSON after validation so unknown provider extensions remain intact. +const StreamEvent = z.object({ + id: z.string().nullish(), + created: z.number().nullish(), + model: z.string().nullish(), + choices: z.array( + z.object({ + delta: z + .object({ + role: z.enum(['assistant', '']).nullish(), + content: z.string().nullish(), + reasoning_content: z.string().nullish(), + reasoning: z.string().nullish(), + tool_calls: z + .array( + z.object({ + index: z.number().nullish(), + id: z.string().nullish(), + function: z.object({ name: z.string().nullish(), arguments: z.string().nullish() }), + extra_content: z + .object({ + google: z.object({ thought_signature: z.string().nullish() }).nullish(), + }) + .nullish(), + }) + ) + .nullish(), + }) + .nullish(), + finish_reason: z.string().nullish(), + }) + ), + usage: z + .object({ + prompt_tokens: z.number().nullish(), + completion_tokens: z.number().nullish(), + total_tokens: z.number().nullish(), + prompt_tokens_details: z.object({ cached_tokens: z.number().nullish() }).nullish(), + completion_tokens_details: z + .object({ + reasoning_tokens: z.number().nullish(), + accepted_prediction_tokens: z.number().nullish(), + rejected_prediction_tokens: z.number().nullish(), + }) + .nullish(), + }) + .nullish(), +}); +const StreamError = z.object({ code: z.coerce.number().int().min(400).max(599) }); +export const mediaType = (headers: Headers) => + headers.get('content-type')?.split(';')[0].trim().toLowerCase(); + +type StreamContext = { + signal: AbortSignal; + abort: AbortController; + // Keep the caller's sanitized error contract without importing its authority dependencies. + failure: (status: number) => unknown; + errorStatus: (error: unknown, invalid: number) => number; +}; +export async function streamHarnessModel( + response: Response, + headers: Headers, + { signal, abort, failure, errorStatus }: StreamContext +) { + const requestId = response.headers.get('request-id'); + if (requestId) headers.set('request-id', requestId); + if (!response.ok || mediaType(response.headers) !== 'text/event-stream' || !response.body) { + await response.body?.cancel().catch(() => undefined); + const status = response.ok ? 422 : response.status; + return Response.json(failure(status), { status, headers }); + } + const reader = boundedBody(response.body, signal) + .pipeThrough(new TextDecoderStream('utf-8', { fatal: true })) + .pipeThrough(new EventSourceParserStream()) + .getReader(); + const encoder = new TextEncoder(); + let cancelled = false; + const close = async () => { + abort.abort(); + await reader.cancel().catch(() => undefined); + reader.releaseLock(); + }; + headers.set('content-type', 'text/event-stream'); + headers.set('content-encoding', 'identity'); + return new Response( + new ReadableStream({ + async pull(controller) { + try { + signal.throwIfAborted(); + const { done, value } = await reader.read(); + if (cancelled) return; + if (done) { + controller.close(); + await close(); + return; + } + let data = value.data; + let terminal = data === '[DONE]'; + if (!terminal) { + const json: unknown = JSON.parse(data); + const parsed = StreamEnvelope.parse(json); + if (parsed.error != null || value.event === 'error') { + const error = StreamError.safeParse(parsed.error); + data = JSON.stringify(failure(error.success ? error.data.code : 422)); + terminal = true; + } else { + StreamEvent.parse(json); + data = JSON.stringify(json); + } + } + controller.enqueue(encoder.encode(`data: ${data}\n\n`)); + if (terminal) { + controller.close(); + await close(); + } + } catch (error) { + if (!cancelled) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(failure(errorStatus(error, 422)))}\n\n`) + ); + controller.close(); + } + await close(); + } + }, + async cancel() { + cancelled = true; + await close(); + }, + }), + { headers } + ); +}