103 lines
9.2 KiB
TypeScript
103 lines
9.2 KiB
TypeScript
import assert from 'node:assert/strict'
|
|
import { createServer } from 'node:http'
|
|
import { once } from 'node:events'
|
|
import { Gateway, type GatewayDelta, type WireMessage } from '../src/main/gateway'
|
|
import { buildAnthropicRequest, buildGeminiRequest } from '../src/main/gateway/native-protocols'
|
|
import { discoverModelConnection } from '../src/main/model-connection'
|
|
import type { GatewayConfig } from '../src/shared/types'
|
|
|
|
const checks: string[] = []
|
|
let mode = 'tools', protocol: 'anthropic' | 'gemini' = 'anthropic', noKey = false
|
|
const observed: Array<{ path: string; body: any }> = []
|
|
const sse = (rows: unknown[]): string => rows.map(row => `data: ${JSON.stringify(row)}\n\n`).join('')
|
|
function anthropic(): unknown[] {
|
|
const rows: any[] = [
|
|
{ type: 'message_start', message: { usage: { input_tokens: 10, output_tokens: 0, cache_read_input_tokens: 4, cache_creation_input_tokens: 2 } } },
|
|
{ type: 'content_block_start', index: 0, content_block: { type: 'thinking', thinking: '' } },
|
|
{ type: 'content_block_delta', index: 0, delta: { type: 'thinking_delta', thinking: '考虑步骤' } },
|
|
{ type: 'content_block_delta', index: 0, delta: { type: 'signature_delta', signature: 'controlled-signature' } },
|
|
{ type: 'content_block_stop', index: 0 },
|
|
{ type: 'content_block_start', index: 1, content_block: { type: 'tool_use', id: 'native-call', name: 'inspect', input: {} } },
|
|
{ type: 'content_block_delta', index: 1, delta: { type: 'input_json_delta', partial_json: '{"path":' } },
|
|
{ type: 'content_block_delta', index: 1, delta: { type: 'input_json_delta', partial_json: '"src"}' } },
|
|
{ type: 'content_block_stop', index: 1 }
|
|
]
|
|
if (mode !== 'truncated') rows.push({ type: 'message_delta', delta: { stop_reason: mode === 'limit' ? 'max_tokens' : 'tool_use' }, usage: { output_tokens: 8 } }, { type: 'message_stop' })
|
|
if (mode === 'error') rows.push({ type: 'error', error: { message: 'private-key-that-must-not-leak' } })
|
|
return rows
|
|
}
|
|
function gemini(): unknown[] {
|
|
const rows: any[] = [{ candidates: [{ index: 0, content: { role: 'model', parts: [{ thought: true, text: '考虑步骤' }] } }] },
|
|
{ candidates: [{ index: 0, content: { role: 'model', parts: [{ functionCall: { id: 'native-call', name: 'inspect', args: { path: 'src' } }, thoughtSignature: 'controlled-signature' }] },
|
|
...(mode !== 'truncated' ? { finishReason: mode === 'limit' ? 'MAX_TOKENS' : 'STOP' } : {}) }],
|
|
usageMetadata: { promptTokenCount: 16, candidatesTokenCount: 5, thoughtsTokenCount: 3, totalTokenCount: 24, cachedContentTokenCount: 4 } }]
|
|
if (mode === 'error') rows.push({ error: { message: 'private-key-that-must-not-leak' } })
|
|
return rows
|
|
}
|
|
const server = createServer(async (req, res) => {
|
|
const url = new URL(req.url!, 'http://local')
|
|
if (url.pathname.startsWith('/root-test')) {
|
|
assert.equal(url.pathname, '/root-test/chat/completions')
|
|
let bytes = ''; for await (const chunk of req) bytes += chunk
|
|
const body = JSON.parse(bytes)
|
|
assert.equal(body.max_completion_tokens, 1024); assert.equal(body.max_tokens, undefined)
|
|
assert.equal(body.reasoning_effort, 'none'); assert.equal(body.temperature, undefined)
|
|
res.writeHead(200, { 'content-type': 'text/event-stream' }); res.end(sse([{ choices: [{ delta: { content: 'OK' }, finish_reason: 'stop' }] }]) + 'data: [DONE]\n\n'); return
|
|
}
|
|
const authenticated = noKey ? !req.headers.authorization && !req.headers['x-api-key'] && !req.headers['x-goog-api-key']
|
|
: protocol === 'anthropic' ? req.headers['x-api-key'] === 'controlled-key' : req.headers['x-goog-api-key'] === 'controlled-key'
|
|
if (!authenticated) { res.writeHead(401); res.end('{}'); return }
|
|
if (req.method === 'GET') {
|
|
res.setHeader('content-type', 'application/json')
|
|
res.end(JSON.stringify(protocol === 'anthropic'
|
|
? url.searchParams.has('after_id') ? { data: [{ id: 'second' }, { id: 'first' }], has_more: false } : { data: [{ id: 'first' }], has_more: true, last_id: 'first' }
|
|
: url.searchParams.has('pageToken') ? { models: [{ name: 'models/second', supportedGenerationMethods: ['generateContent'] }, { name: 'models/embedding', supportedGenerationMethods: ['embedContent'] }] }
|
|
: { models: [{ name: 'models/first', supportedGenerationMethods: ['generateContent'] }], nextPageToken: 'next' })); return
|
|
}
|
|
let bytes = ''; for await (const chunk of req) bytes += chunk
|
|
observed.push({ path: req.url!, body: JSON.parse(bytes) })
|
|
res.writeHead(200, { 'content-type': 'text/event-stream' }); res.end(sse(protocol === 'anthropic' ? anthropic() : gemini()))
|
|
})
|
|
server.listen(0, '127.0.0.1'); await once(server, 'listening')
|
|
const address = server.address(); assert(address && typeof address === 'object')
|
|
const root = `http://127.0.0.1:${address.port}`
|
|
const collect = async (config: GatewayConfig): Promise<GatewayDelta[]> => { const events: GatewayDelta[] = []; for await (const event of new Gateway(config).stream({ messages: [{ role: 'user', content: 'inspect src' }] })) events.push(event); return events }
|
|
try {
|
|
const geminiAuto: GatewayConfig = { api: 'google-generative-ai', baseUrl: root, model: 'gemini-3.1-pro-preview', apiKey: 'controlled-key', temperature: 0.5, maxToolRounds: 3,
|
|
supportsReasoning: true, thinkingLevel: 'auto', reasoningOptions: [{ type: 'effort', values: ['low', 'medium', 'high'] }] }
|
|
assert.deepEqual(buildGeminiRequest(geminiAuto, { messages: [] }).generationConfig.thinkingConfig, { includeThoughts: true })
|
|
assert.deepEqual(buildGeminiRequest({ ...geminiAuto, thinkingLevel: 'low' }, { messages: [] }).generationConfig.thinkingConfig, { thinkingLevel: 'LOW', includeThoughts: true })
|
|
checks.push('Gemini effort-only models use native thinking levels without inventing a token budget for auto')
|
|
const rootConfig: GatewayConfig = { api: 'openai-completions', baseUrl: root + '/root-test', model: 'gpt-5.4', apiKey: 'controlled-key', temperature: 0.5, maxToolRounds: 3,
|
|
maxOutputTokens: 1024, supportsReasoning: true, thinkingLevel: 'off', reasoningOptions: [{ type: 'effort', values: ['none', 'low', 'high'] }], temperatureSupported: false }
|
|
assert((await collect(rootConfig)).some(event => event.type === 'text' && event.text === 'OK'))
|
|
checks.push('OpenAI fallback honors the negotiated root, modern completion budget and disabled reasoning/temperature')
|
|
for (protocol of ['anthropic', 'gemini'] as const) {
|
|
noKey = false; mode = 'tools'
|
|
const config: GatewayConfig = { api: protocol === 'anthropic' ? 'anthropic-messages' : 'google-generative-ai', baseUrl: root + (protocol === 'anthropic' ? '/v1' : '/v1beta'), model: 'controlled-model', apiKey: 'controlled-key', temperature: 0, maxToolRounds: 3 }
|
|
const discovery = await discoverModelConnection(root, protocol, 'controlled-key', true)
|
|
assert.deepEqual(discovery.models.map(model => model.id), ['first', 'second']); checks.push(protocol + ' lists every page, filters non-generative IDs and removes duplicates')
|
|
const events = await collect(config), native = events.find(event => event.type === 'native_content')!, calls = events.find(event => event.type === 'tool_calls')!
|
|
assert(native.type === 'native_content' && calls.type === 'tool_calls')
|
|
assert.deepEqual(calls.calls[0], { id: 'native-call', type: 'function', function: { name: 'inspect', arguments: '{"path":"src"}' } })
|
|
assert.equal(events.filter(event => event.type === 'reasoning').map(event => event.text).join(''), '考虑步骤')
|
|
const usage = events.filter(event => event.type === 'usage').at(-1)!
|
|
assert(usage.type === 'usage'); assert.deepEqual(usage.usage, { inputTokens: 16, outputTokens: 8, totalTokens: 24, cachedTokens: 4 })
|
|
const messages: WireMessage[] = [{ role: 'user', content: 'inspect src' }, { role: 'assistant', content: null, tool_calls: calls.calls, providerContent: { api: native.api, content: native.content } }, { role: 'tool', tool_call_id: 'native-call', content: 'contents' }]
|
|
const body = protocol === 'anthropic' ? buildAnthropicRequest(config, { messages }) : buildGeminiRequest(config, { messages })
|
|
assert(JSON.stringify(body).includes('controlled-signature'))
|
|
assert(JSON.stringify(body).includes('native-call')); assert(JSON.stringify(body).includes('contents'))
|
|
checks.push(protocol + ' native tool IDs, arguments, reasoning, signatures and usage survive tool continuation')
|
|
for (mode of ['truncated', 'limit', 'error']) {
|
|
const pending: GatewayDelta[] = []
|
|
await assert.rejects(async () => { for await (const event of new Gateway(config).stream({ messages: [{ role: 'user', content: 'go' }] })) pending.push(event) }, error => error instanceof Error && !error.message.includes('private-key-that-must-not-leak'))
|
|
assert(!pending.some(event => event.type === 'tool_calls')); checks.push(protocol + ' ' + mode + ' never dispatches incomplete tools or leaks provider error content')
|
|
}
|
|
mode = 'tools'; noKey = true
|
|
assert.equal((await discoverModelConnection(root, protocol, '', true)).models.length, 2)
|
|
assert((await collect({ ...config, authMode: 'none', apiKey: '' })).some(event => event.type === 'tool_calls'))
|
|
checks.push(protocol + ' keyless custom endpoints send no authentication headers')
|
|
}
|
|
console.log('CUSTOM_MODEL_NATIVE_PASS ' + JSON.stringify({ checks }))
|
|
} finally { server.closeAllConnections(); await new Promise<void>(resolve => server.close(() => resolve())) }
|