Files
ycode/scripts/chat-completions-usage-regression.ts
T

237 lines
8.7 KiB
TypeScript

import assert from 'node:assert/strict'
import { createServer } from 'node:http'
import { Gateway, GatewayError, type GatewayDelta } from '../src/main/gateway'
const requests = new Map<string, Array<Record<string, unknown>>>()
const abortConnections = new Set<import('node:http').ServerResponse>()
function attempts(model: string): Array<Record<string, unknown>> {
return requests.get(model) ?? []
}
function recordAttempt(body: Record<string, unknown>): number {
const model = String(body.model)
const items = requests.get(model) ?? []
items.push(body)
requests.set(model, items)
return items.length
}
function successSse(response: import('node:http').ServerResponse, usage: Record<string, unknown>): void {
response.writeHead(200, { 'content-type': 'text/event-stream' })
response.end([
{ choices: [{ delta: { content: 'ok' }, finish_reason: null }] },
{ choices: [{ delta: {}, finish_reason: 'stop' }] },
{ choices: [], usage }
].map((frame) => `data: ${JSON.stringify(frame)}\n\n`).join('') + 'data: [DONE]\n\n')
}
const server = createServer(async (request, response) => {
const chunks: Buffer[] = []
for await (const chunk of request) chunks.push(Buffer.from(chunk))
const body = JSON.parse(Buffer.concat(chunks).toString()) as Record<string, unknown>
const model = String(body.model)
const attempt = recordAttempt(body)
const validUsage = {
prompt_tokens: 128,
completion_tokens: 2,
total_tokens: 130,
prompt_tokens_details: { cached_tokens: 32 }
}
if (model === 'valid' || model === 'invalid-usage') {
successSse(response, model === 'invalid-usage'
? { prompt_tokens: -1, completion_tokens: 2, total_tokens: 1 }
: validUsage)
return
}
if (model === 'legacy-400') {
if (attempt === 1 && body.stream_options) {
response.writeHead(400, { 'content-type': 'application/json' })
response.end(JSON.stringify({ error: { message: 'Unsupported parameter: stream_options' } }))
return
}
successSse(response, validUsage)
return
}
if (model === 'legacy-422') {
if (attempt === 1 && body.stream_options) {
response.writeHead(422, { 'content-type': 'application/json' })
response.end(JSON.stringify({ error: { message: 'include_usage is not supported by this server' } }))
return
}
successSse(response, validUsage)
return
}
if (model === 'ordinary-400') {
response.writeHead(400, { 'content-type': 'application/json' })
response.end(JSON.stringify({
error: { message: 'request validation failed while parsing stream_options' },
apiKey: 'UPSTREAM_SECRET_MUST_NOT_ECHO'
}))
return
}
if (model === 'auth-401') {
response.writeHead(401, { 'content-type': 'application/json' })
response.end(JSON.stringify({
error: { message: 'invalid api key' },
authorization: 'Bearer UPSTREAM_SECRET_MUST_NOT_ECHO'
}))
return
}
if (model === 'server-500-options') {
response.writeHead(500, { 'content-type': 'application/json' })
response.end(JSON.stringify({ error: { message: 'stream_options unsupported due to server failure' } }))
return
}
if (model === 'started-sse-error') {
response.writeHead(200, { 'content-type': 'text/event-stream' })
response.write(`data: ${JSON.stringify({ choices: [{ delta: { content: 'before-error' }, finish_reason: null }] })}\n\n`)
response.write(`data: ${JSON.stringify({ error: { message: 'stream_options unsupported' } })}\n\n`)
response.end()
return
}
if (model === 'abort-retry') {
if (attempt === 1 && body.stream_options) {
response.writeHead(400, { 'content-type': 'application/json' })
response.end(JSON.stringify({ error: { message: 'unknown field: stream_options' } }))
return
}
abortConnections.add(response)
response.writeHead(200, {
'content-type': 'text/event-stream',
'cache-control': 'no-cache',
connection: 'keep-alive'
})
response.flushHeaders()
const keepalive = setInterval(() => {
if (!response.destroyed) response.write(': keepalive\n\n')
}, 100)
response.on('close', () => {
clearInterval(keepalive)
abortConnections.delete(response)
})
return
}
response.writeHead(400, { 'content-type': 'application/json' })
response.end(JSON.stringify({ error: { message: `unknown QA model: ${model}` } }))
})
function gatewayFor(model: string): Gateway {
const address = server.address()
assert.ok(address && typeof address === 'object')
return new Gateway({
baseUrl: `http://127.0.0.1:${address.port}/v1`,
apiKey: '',
authMode: 'none',
model,
temperature: 0,
maxToolRounds: 1
})
}
async function collect(model: string, signal?: AbortSignal): Promise<GatewayDelta[]> {
const events: GatewayDelta[] = []
for await (const event of gatewayFor(model).stream({
messages: [{ role: 'user', content: 'test' }],
signal
})) {
events.push(event)
}
return events
}
async function captureError(model: string): Promise<GatewayError> {
try {
await collect(model)
} catch (error) {
assert.ok(error instanceof GatewayError, `${model} must reject with GatewayError`)
return error
}
assert.fail(`${model} should reject`)
}
async function waitForAttempts(model: string, count: number): Promise<void> {
for (let index = 0; index < 250; index++) {
if (attempts(model).length === count) return
await new Promise((resolve) => setTimeout(resolve, 10))
}
assert.fail(`${model} did not reach ${count} requests`)
}
await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
try {
const valid = await collect('valid')
assert.deepEqual(valid.filter((event) => event.type === 'text'), [{ type: 'text', text: 'ok' }])
assert.deepEqual(valid.filter((event) => event.type === 'finish'), [{ type: 'finish', reason: 'stop' }])
assert.deepEqual(valid.filter((event) => event.type === 'usage'), [{
type: 'usage',
usage: { inputTokens: 128, outputTokens: 2, totalTokens: 130, cachedTokens: 32 }
}])
assert.equal(attempts('valid').length, 1)
assert.deepEqual(attempts('valid')[0]?.stream_options, { include_usage: true })
const invalidUsage = await collect('invalid-usage')
assert.equal(invalidUsage.filter((event) => event.type === 'usage').length, 0,
'invalid usage must not emit a usage event')
for (const model of ['legacy-400', 'legacy-422']) {
const events = await collect(model)
assert.deepEqual(events.filter((event) => event.type === 'text'), [{ type: 'text', text: 'ok' }])
assert.deepEqual(events.filter((event) => event.type === 'usage'), [{
type: 'usage',
usage: { inputTokens: 128, outputTokens: 2, totalTokens: 130, cachedTokens: 32 }
}])
const tries = attempts(model)
assert.equal(tries.length, 2, `${model} must retry exactly once`)
assert.deepEqual(tries[0]?.stream_options, { include_usage: true })
assert.equal(Object.hasOwn(tries[1] ?? {}, 'stream_options'), false,
`${model} retry must omit stream_options`)
}
const ordinary = await captureError('ordinary-400')
assert.equal(ordinary.status, 400)
assert.equal(attempts('ordinary-400').length, 1, 'an ordinary 400 must not retry')
assert.equal(ordinary.message.includes('UPSTREAM_SECRET_MUST_NOT_ECHO'), false,
'upstream body content must not be copied into the renderer-facing message')
const auth = await captureError('auth-401')
assert.equal(auth.status, 401)
assert.equal(attempts('auth-401').length, 1, 'authentication failures must not retry')
assert.equal(auth.message.includes('UPSTREAM_SECRET_MUST_NOT_ECHO'), false,
'authorization material must not be copied into the renderer-facing message')
const serverFailure = await captureError('server-500-options')
assert.equal(serverFailure.status, 500)
assert.equal(attempts('server-500-options').length, 1,
'a 5xx response must not retry even when its body names stream_options')
const startedSse = await captureError('started-sse-error')
assert.equal(attempts('started-sse-error').length, 1,
'an error after SSE has started must not trigger another request')
assert.match(startedSse.message, /上游错误/)
const controller = new AbortController()
const aborted = collect('abort-retry', controller.signal)
await waitForAttempts('abort-retry', 2)
controller.abort()
await assert.rejects(aborted, (error: unknown) =>
error instanceof Error && error.name === 'AbortError')
assert.equal(attempts('abort-retry').length, 2,
'the shared AbortSignal must cancel the in-flight compatibility retry')
console.log('chat-completions-usage-regression: all checks passed')
} finally {
for (const response of abortConnections) response.destroy()
server.closeAllConnections?.()
await new Promise<void>((resolve) => server.close(() => resolve()))
}