From da5d9b31e952049c80be373ccec14f31b6984083 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:42:56 +0800 Subject: [PATCH 001/157] feat(gateway): add encrypted affinity blob codec Encode gateway affinity metadata as an authenticated AES-GCM trailer with a two-byte length marker. Preserve canonical base64 and base64url payload bytes without nested encoding, classify other strings as raw UTF-8, and treat unrecognized or unauthenticated carriers as foreign for gateway cascades. --- .../src/data-plane/chat/affinity/codec.ts | 187 ++++++++++++++++++ .../data-plane/chat/affinity/codec_test.ts | 92 +++++++++ .../src/data-plane/chat/affinity/types.ts | 23 +++ 3 files changed, 302 insertions(+) create mode 100644 packages/gateway/src/data-plane/chat/affinity/codec.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/codec_test.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/types.ts diff --git a/packages/gateway/src/data-plane/chat/affinity/codec.ts b/packages/gateway/src/data-plane/chat/affinity/codec.ts new file mode 100644 index 000000000..c34f3bb70 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/codec.ts @@ -0,0 +1,187 @@ +import type { AffinityEnvelope, AffinityOrigin, AffinityTarget, DecodedAffinityBlob } from './types.ts'; +import type { AliasRules } from '@floway-dev/protocols/common'; + +const IV_BYTES = 12; +const LENGTH_BYTES = 2; +const MAX_ENCRYPTED_BYTES = 0xffff; +const textEncoder = new TextEncoder(); +const fatalTextDecoder = new TextDecoder('utf-8', { fatal: true }); + +const bytesToBase64 = (bytes: Uint8Array): string => { + let binary = ''; + for (let offset = 0; offset < bytes.length; offset += 0x8000) { + binary += String.fromCharCode(...bytes.subarray(offset, offset + 0x8000)); + } + return btoa(binary); +}; + +const base64ToBytes = (value: string): Uint8Array => { + const binary = atob(value); + const bytes = new Uint8Array(binary.length); + for (let index = 0; index < binary.length; index += 1) bytes[index] = binary.charCodeAt(index); + return bytes; +}; + +const bytesToBase64url = (bytes: Uint8Array): string => + bytesToBase64(bytes).replaceAll('+', '-').replaceAll('/', '_').replace(/=+$/, ''); + +const base64urlToBytes = (value: string): Uint8Array => { + const standard = value.replaceAll('-', '+').replaceAll('_', '/'); + const padding = (4 - standard.length % 4) % 4; + return base64ToBytes(`${standard}${'='.repeat(padding)}`); +}; + +const decodeCanonicalBase64 = (value: string): Uint8Array | null => { + try { + const bytes = base64ToBytes(value); + return bytesToBase64(bytes) === value ? bytes : null; + } catch { + return null; + } +}; + +const decodeCanonicalBase64url = (value: string): Uint8Array | null => { + try { + const bytes = base64urlToBytes(value); + return bytesToBase64url(bytes) === value ? bytes : null; + } catch { + return null; + } +}; + +const decodeOriginal = (value: string): { bytes: Uint8Array; origin: AffinityOrigin } => { + if (value.length > 0) { + const base64 = decodeCanonicalBase64(value); + if (base64 !== null) return { bytes: base64, origin: 'base64' }; + const base64url = decodeCanonicalBase64url(value); + if (base64url !== null) return { bytes: base64url, origin: 'base64url' }; + } + return { bytes: textEncoder.encode(value), origin: 'raw' }; +}; + +const encodeOriginal = (bytes: Uint8Array, origin: AffinityOrigin): string => { + switch (origin) { + case 'base64': return bytesToBase64(bytes); + case 'base64url': return bytesToBase64url(bytes); + case 'raw': return fatalTextDecoder.decode(bytes); + } +}; + +const parseHexSecret = (secret: string): Uint8Array => { + if (!/^[0-9a-f]{64}$/.test(secret)) throw new TypeError('Affinity secret must be 64 lowercase hexadecimal characters'); + const bytes = new Uint8Array(32); + for (let index = 0; index < bytes.length; index += 1) { + bytes[index] = Number.parseInt(secret.slice(index * 2, index * 2 + 2), 16); + } + return bytes; +}; + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null && !Array.isArray(value); + +const parseEnvelope = (value: unknown): AffinityEnvelope | null => { + if (!isRecord(value) || value.version !== 1 || !isRecord(value.affinity)) return null; + const origin = value.origin; + if (origin !== undefined && origin !== 'raw' && origin !== 'base64' && origin !== 'base64url') return null; + + const affinity = value.affinity; + if ( + (affinity.mode !== 'prefer' && affinity.mode !== 'force') + || typeof affinity.upstreamId !== 'string' + || typeof affinity.upstreamRevision !== 'string' + || typeof affinity.modelId !== 'string' + || typeof affinity.rulesPresent !== 'boolean' + || (affinity.upstreamItemId !== undefined && typeof affinity.upstreamItemId !== 'string') + || (affinity.rulesPresent && !isRecord(affinity.rules)) + || (!affinity.rulesPresent && affinity.rules !== undefined) + ) return null; + + const parsedAffinity: AffinityTarget = { + mode: affinity.mode, + upstreamId: affinity.upstreamId, + upstreamRevision: affinity.upstreamRevision, + modelId: affinity.modelId, + rulesPresent: affinity.rulesPresent, + ...(affinity.rulesPresent ? { rules: affinity.rules as AliasRules } : {}), + ...(affinity.upstreamItemId !== undefined ? { upstreamItemId: affinity.upstreamItemId } : {}), + }; + return { + version: 1, + ...(origin !== undefined ? { origin } : {}), + affinity: parsedAffinity, + }; +}; + +const concatBytes = (...parts: readonly Uint8Array[]): Uint8Array => { + const result = new Uint8Array(parts.reduce((sum, part) => sum + part.length, 0)); + let offset = 0; + for (const part of parts) { + result.set(part, offset); + offset += part.length; + } + return result; +}; + +const encryptedLengthMarker = (length: number): Uint8Array => + new Uint8Array([length >>> 8, length & 0xff]); + +const encryptedLengthFrom = (bytes: Uint8Array): number => + (bytes[bytes.length - 2] << 8) | bytes[bytes.length - 1]; + +export class AffinityCodec { + readonly #key: Promise; + + constructor(secret: string) { + this.#key = crypto.subtle.importKey('raw', parseHexSecret(secret), 'AES-GCM', false, ['encrypt', 'decrypt']); + } + + async wrap(value: string | undefined, affinity: AffinityTarget): Promise { + const original = value === undefined ? undefined : decodeOriginal(value); + const envelope: AffinityEnvelope = { + version: 1, + ...(original !== undefined ? { origin: original.origin } : {}), + affinity, + }; + const iv = crypto.getRandomValues(new Uint8Array(IV_BYTES)); + const ciphertext = new Uint8Array(await crypto.subtle.encrypt( + { name: 'AES-GCM', iv }, + await this.#key, + textEncoder.encode(JSON.stringify(envelope)), + )); + const encrypted = concatBytes(iv, ciphertext); + if (encrypted.length > MAX_ENCRYPTED_BYTES) throw new RangeError('Encrypted affinity envelope exceeds the 2-byte length marker'); + const framed = concatBytes(original?.bytes ?? new Uint8Array(), encrypted, encryptedLengthMarker(encrypted.length)); + return original?.origin === 'base64url' ? bytesToBase64url(framed) : bytesToBase64(framed); + } + + async unwrap(value: string): Promise { + const framed = decodeCanonicalBase64(value) ?? decodeCanonicalBase64url(value); + if (framed === null || framed.length < LENGTH_BYTES + IV_BYTES + 16) return { kind: 'foreign', value }; + + const encryptedLength = encryptedLengthFrom(framed); + const originalLength = framed.length - LENGTH_BYTES - encryptedLength; + if (encryptedLength < IV_BYTES + 16 || originalLength < 0) return { kind: 'foreign', value }; + + const encrypted = framed.subarray(originalLength, framed.length - LENGTH_BYTES); + const iv = encrypted.subarray(0, IV_BYTES); + const ciphertext = encrypted.subarray(IV_BYTES); + try { + const plaintext = await crypto.subtle.decrypt( + { name: 'AES-GCM', iv }, + await this.#key, + ciphertext, + ); + const envelope = parseEnvelope(JSON.parse(fatalTextDecoder.decode(plaintext)) as unknown); + if (envelope === null) return { kind: 'foreign', value }; + const original = framed.subarray(0, originalLength); + if (envelope.origin === undefined) { + return original.length === 0 + ? { kind: 'owned', envelope } + : { kind: 'foreign', value }; + } + return { kind: 'owned', value: encodeOriginal(original, envelope.origin), envelope }; + } catch { + return { kind: 'foreign', value }; + } + } +} diff --git a/packages/gateway/src/data-plane/chat/affinity/codec_test.ts b/packages/gateway/src/data-plane/chat/affinity/codec_test.ts new file mode 100644 index 000000000..783e4ea45 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/codec_test.ts @@ -0,0 +1,92 @@ +import { describe, expect, test } from 'vitest'; + +import { AffinityCodec } from './codec.ts'; +import type { AffinityTarget } from './types.ts'; + +const SECRET = '00'.repeat(32); +const OTHER_SECRET = '11'.repeat(32); +const affinity: AffinityTarget = { + mode: 'prefer', + upstreamId: 'upstream-a', + upstreamRevision: '2026-07-15T00:00:00.000Z', + modelId: 'model-a', + rulesPresent: false, +}; + +describe('AffinityCodec', () => { + test.each([ + ['raw', 'not base64!'], + ['base64', btoa('upstream opaque bytes')], + ['base64url', '--__'], + ['empty raw', ''], + ])('round-trips %s input', async (_label, original) => { + const codec = new AffinityCodec(SECRET); + const wrapped = await codec.wrap(original, affinity); + const decoded = await codec.unwrap(wrapped); + + expect(decoded).toEqual({ + kind: 'owned', + value: original, + envelope: { + version: 1, + origin: _label === 'base64' ? 'base64' : _label === 'base64url' ? 'base64url' : 'raw', + affinity, + }, + }); + }); + + test('uses no origin for a synthetic carrier', async () => { + const codec = new AffinityCodec(SECRET); + const wrapped = await codec.wrap(undefined, affinity); + + expect(await codec.unwrap(wrapped)).toEqual({ + kind: 'owned', + envelope: { version: 1, affinity }, + }); + }); + + test('does not base64-encode canonical base64 bytes a second time', async () => { + const codec = new AffinityCodec(SECRET); + const originalBytes = crypto.getRandomValues(new Uint8Array(48)); + const original = btoa(String.fromCharCode(...originalBytes)); + const wrapped = await codec.wrap(original, affinity); + const framedBytes = Uint8Array.from(atob(wrapped), char => char.charCodeAt(0)); + + expect(framedBytes.subarray(0, originalBytes.length)).toEqual(originalBytes); + }); + + test('preserves a foreign value byte-for-byte on authentication failure', async () => { + const wrapped = await new AffinityCodec(SECRET).wrap('opaque', affinity); + + expect(await new AffinityCodec(OTHER_SECRET).unwrap(wrapped)).toEqual({ kind: 'foreign', value: wrapped }); + }); + + test('preserves malformed and tampered values as foreign', async () => { + const codec = new AffinityCodec(SECRET); + const wrapped = await codec.wrap('opaque', affinity); + const bytes = Uint8Array.from(atob(wrapped), char => char.charCodeAt(0)); + bytes[bytes.length - 3] ^= 1; + const tampered = btoa(String.fromCharCode(...bytes)); + + expect(await codec.unwrap('not-an-envelope')).toEqual({ kind: 'foreign', value: 'not-an-envelope' }); + expect(await codec.unwrap(tampered)).toEqual({ kind: 'foreign', value: tampered }); + }); + + test('unwraps nested gateway carriers one layer at a time', async () => { + const innerCodec = new AffinityCodec(OTHER_SECRET); + const outerCodec = new AffinityCodec(SECRET); + const inner = await innerCodec.wrap('upstream', affinity); + const outer = await outerCodec.wrap(inner, { ...affinity, upstreamId: 'inner-gateway' }); + + const outerDecoded = await outerCodec.unwrap(outer); + expect(outerDecoded.kind).toBe('owned'); + if (outerDecoded.kind !== 'owned') throw new Error('Expected owned outer carrier'); + expect(outerDecoded.value).toBe(inner); + expect(await innerCodec.unwrap(outerDecoded.value!)).toMatchObject({ kind: 'owned', value: 'upstream' }); + }); + + test('rejects malformed secrets', () => { + expect(() => new AffinityCodec('00')).toThrow(TypeError); + expect(() => new AffinityCodec('AA'.repeat(32))).toThrow(TypeError); + }); +}); diff --git a/packages/gateway/src/data-plane/chat/affinity/types.ts b/packages/gateway/src/data-plane/chat/affinity/types.ts new file mode 100644 index 000000000..12dba6554 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/types.ts @@ -0,0 +1,23 @@ +import type { AliasRules } from '@floway-dev/protocols/common'; + +export type AffinityOrigin = 'raw' | 'base64' | 'base64url'; + +export interface AffinityTarget { + mode: 'prefer' | 'force'; + upstreamId: string; + upstreamRevision: string; + modelId: string; + rulesPresent: boolean; + rules?: AliasRules; + upstreamItemId?: string; +} + +export interface AffinityEnvelope { + version: 1; + origin?: AffinityOrigin; + affinity: AffinityTarget; +} + +export type DecodedAffinityBlob = + | { kind: 'foreign'; value: string } + | { kind: 'owned'; value?: string; envelope: AffinityEnvelope }; From eaba9587da3675caa35737b04d594bde48074908 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:47:16 +0800 Subject: [PATCH 002/157] feat(gateway): identify exact affinity targets Carry the upstream configuration revision on provider candidates and encode direct-versus-alias rule identity explicitly. Route forcing affinity only to an exact target and use the latest available preferred target without narrowing normal fallback order. --- .../src/data-plane/chat/affinity/candidate.ts | 76 +++++++++++++++++++ .../chat/affinity/candidate_test.ts | 60 +++++++++++++++ packages/provider-azure/src/provider.ts | 1 + packages/provider-claude-code/src/provider.ts | 1 + packages/provider-codex/src/provider.ts | 1 + packages/provider-copilot/src/provider.ts | 1 + packages/provider-custom/src/provider.ts | 1 + packages/provider-ollama/src/provider.ts | 1 + packages/provider/src/provider.ts | 4 + packages/test-utils/src/stubs.ts | 1 + 10 files changed, 147 insertions(+) create mode 100644 packages/gateway/src/data-plane/chat/affinity/candidate.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/candidate_test.ts diff --git a/packages/gateway/src/data-plane/chat/affinity/candidate.ts b/packages/gateway/src/data-plane/chat/affinity/candidate.ts new file mode 100644 index 000000000..64f6ffb50 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/candidate.ts @@ -0,0 +1,76 @@ +import { isEqual } from 'es-toolkit'; + +import type { AffinityTarget } from './types.ts'; +import type { RoutingDecision } from '../shared/routing.ts'; +import type { ModelCandidate } from '@floway-dev/provider'; + +const sameTarget = (left: AffinityTarget, right: AffinityTarget): boolean => + left.upstreamId === right.upstreamId + && left.upstreamRevision === right.upstreamRevision + && left.modelId === right.modelId + && left.rulesPresent === right.rulesPresent + && isEqual(left.rules, right.rules); + +export const affinityTargetForCandidate = ( + candidate: ModelCandidate, + mode: AffinityTarget['mode'], + upstreamItemId?: string, +): AffinityTarget => ({ + mode, + upstreamId: candidate.provider.upstream, + upstreamRevision: candidate.provider.upstreamRevision, + modelId: candidate.model.id, + rulesPresent: candidate.rules !== undefined, + ...(candidate.rules !== undefined ? { rules: candidate.rules } : {}), + ...(upstreamItemId !== undefined ? { upstreamItemId } : {}), +}); + +export const candidateMatchesAffinity = (candidate: ModelCandidate, affinity: AffinityTarget): boolean => + candidate.provider.upstream === affinity.upstreamId + && candidate.provider.upstreamRevision === affinity.upstreamRevision + && candidate.model.id === affinity.modelId + && (candidate.rules !== undefined) === affinity.rulesPresent + && isEqual(candidate.rules, affinity.rules); + +export const routeCandidatesByAffinity = ( + candidates: readonly T[], + affinities: readonly AffinityTarget[], +): RoutingDecision => { + const forcing: AffinityTarget[] = []; + for (const affinity of affinities) { + if (affinity.mode === 'force' && !forcing.some(existing => sameTarget(existing, affinity))) forcing.push(affinity); + } + if (forcing.length > 1) { + return { + kind: 'failure', + failure: { + kind: 'routing-unavailable', + message: `Client-carried state requires multiple incompatible targets: ${forcing.map(target => `'${target.upstreamId}/${target.modelId}'`).join(', ')}.`, + }, + }; + } + if (forcing.length === 1) { + const matching = candidates.filter(candidate => candidateMatchesAffinity(candidate, forcing[0])); + return matching.length > 0 + ? { kind: 'success', candidates: matching } + : { + kind: 'failure', + failure: { + kind: 'routing-unavailable', + message: `Client-carried state requires unavailable target '${forcing[0].upstreamId}/${forcing[0].modelId}'.`, + }, + }; + } + + const preferred = affinities.findLast(affinity => affinity.mode === 'prefer'); + if (preferred === undefined) return { kind: 'success', candidates }; + const matching = candidates.filter(candidate => candidateMatchesAffinity(candidate, preferred)); + if (matching.length === 0) return { kind: 'success', candidates }; + return { + kind: 'success', + candidates: [ + ...matching, + ...candidates.filter(candidate => !candidateMatchesAffinity(candidate, preferred)), + ], + }; +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/candidate_test.ts b/packages/gateway/src/data-plane/chat/affinity/candidate_test.ts new file mode 100644 index 000000000..2c7aa4891 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/candidate_test.ts @@ -0,0 +1,60 @@ +import { describe, expect, test } from 'vitest'; + +import { affinityTargetForCandidate, candidateMatchesAffinity, routeCandidatesByAffinity } from './candidate.ts'; +import type { AliasRules } from '@floway-dev/protocols/common'; +import { stubModelCandidate } from '@floway-dev/test-utils'; + +const candidate = (upstream: string, revision: string, model: string, rules?: AliasRules) => { + const base = stubModelCandidate(); + const value = stubModelCandidate({ + provider: { ...base.provider, upstream, upstreamRevision: revision }, + model: { id: model }, + }); + return rules === undefined ? value : { ...value, rules }; +}; + +describe('client-carried affinity candidate routing', () => { + test('matches upstream revision, model, and the exact alias rules presence', () => { + const direct = candidate('up-a', 'rev-1', 'model-a'); + const alias = candidate('up-a', 'rev-1', 'model-a', {}); + + expect(candidateMatchesAffinity(direct, affinityTargetForCandidate(direct, 'prefer'))).toBe(true); + expect(candidateMatchesAffinity(alias, affinityTargetForCandidate(direct, 'prefer'))).toBe(false); + expect(candidateMatchesAffinity(candidate('up-a', 'rev-2', 'model-a'), affinityTargetForCandidate(direct, 'prefer'))).toBe(false); + }); + + test('moves the latest available preferred target to the front', () => { + const first = candidate('up-a', 'rev', 'model'); + const second = candidate('up-b', 'rev', 'model'); + const decision = routeCandidatesByAffinity( + [first, second], + [affinityTargetForCandidate(first, 'prefer'), affinityTargetForCandidate(second, 'prefer')], + ); + + expect(decision.kind).toBe('success'); + if (decision.kind !== 'success') throw new Error('Expected successful routing'); + expect(decision.candidates).toEqual([second, first]); + }); + + test('keeps normal order when a preferred target is unavailable', () => { + const first = candidate('up-a', 'rev', 'model'); + const second = candidate('up-b', 'rev', 'model'); + const unavailable = candidate('up-c', 'rev', 'model'); + + expect(routeCandidatesByAffinity([first, second], [affinityTargetForCandidate(unavailable, 'prefer')])).toEqual({ + kind: 'success', + candidates: [first, second], + }); + }); + + test('fails unavailable and conflicting force affinity', () => { + const first = candidate('up-a', 'rev', 'model'); + const second = candidate('up-b', 'rev', 'model'); + + expect(routeCandidatesByAffinity([first], [affinityTargetForCandidate(second, 'force')])).toMatchObject({ kind: 'failure' }); + expect(routeCandidatesByAffinity([first, second], [ + affinityTargetForCandidate(first, 'force'), + affinityTargetForCandidate(second, 'force'), + ])).toMatchObject({ kind: 'failure' }); + }); +}); diff --git a/packages/provider-azure/src/provider.ts b/packages/provider-azure/src/provider.ts index 209a0b5ba..d1939f12e 100644 --- a/packages/provider-azure/src/provider.ts +++ b/packages/provider-azure/src/provider.ts @@ -100,6 +100,7 @@ export const createAzureProvider = (record: UpstreamRecord): Provider => { return { upstream: azure.id, + upstreamRevision: azure.updatedAt, kind: 'azure', name: azure.name, disabledPublicModelIds: azure.disabledPublicModelIds, diff --git a/packages/provider-claude-code/src/provider.ts b/packages/provider-claude-code/src/provider.ts index ae4f19437..11bb610b2 100644 --- a/packages/provider-claude-code/src/provider.ts +++ b/packages/provider-claude-code/src/provider.ts @@ -97,6 +97,7 @@ export const createClaudeCodeProvider = (record: UpstreamRecord): Provider => { return { upstream: record.id, + upstreamRevision: record.updatedAt, kind: 'claude-code', name: record.name, disabledPublicModelIds: record.disabledPublicModelIds, diff --git a/packages/provider-codex/src/provider.ts b/packages/provider-codex/src/provider.ts index 995453a91..721f9f812 100644 --- a/packages/provider-codex/src/provider.ts +++ b/packages/provider-codex/src/provider.ts @@ -141,6 +141,7 @@ export const createCodexProvider = (record: UpstreamRecord): Provider => { return { upstream: record.id, + upstreamRevision: record.updatedAt, kind: 'codex', name: record.name, disabledPublicModelIds: record.disabledPublicModelIds, diff --git a/packages/provider-copilot/src/provider.ts b/packages/provider-copilot/src/provider.ts index 800eaea88..de73d4922 100644 --- a/packages/provider-copilot/src/provider.ts +++ b/packages/provider-copilot/src/provider.ts @@ -468,6 +468,7 @@ export const createCopilotProvider = (record: UpstreamRecord): Provider => { return { upstream: copilot.id, + upstreamRevision: copilot.updatedAt, kind: 'copilot', name: copilot.name, disabledPublicModelIds: copilot.disabledPublicModelIds, diff --git a/packages/provider-custom/src/provider.ts b/packages/provider-custom/src/provider.ts index 576894229..cea1343ec 100644 --- a/packages/provider-custom/src/provider.ts +++ b/packages/provider-custom/src/provider.ts @@ -188,6 +188,7 @@ export const createCustomProvider = (record: UpstreamRecord): Provider => { return { upstream: record.id, + upstreamRevision: record.updatedAt, kind: 'custom', name: record.name, disabledPublicModelIds: record.disabledPublicModelIds, diff --git a/packages/provider-ollama/src/provider.ts b/packages/provider-ollama/src/provider.ts index 2cd37a8b0..454a09974 100644 --- a/packages/provider-ollama/src/provider.ts +++ b/packages/provider-ollama/src/provider.ts @@ -177,6 +177,7 @@ export const createOllamaProvider = (record: UpstreamRecord): Provider => { return { upstream: record.id, + upstreamRevision: record.updatedAt, kind: 'ollama', name: record.name, disabledPublicModelIds: record.disabledPublicModelIds, diff --git a/packages/provider/src/provider.ts b/packages/provider/src/provider.ts index c1761524d..01a607fd3 100644 --- a/packages/provider/src/provider.ts +++ b/packages/provider/src/provider.ts @@ -23,6 +23,10 @@ export type ResponsesAction = 'generate' | 'compact'; export interface Provider { upstream: string; + // Changes whenever the operator replaces this upstream's configuration or + // credentials. Opaque reasoning state minted by an earlier revision must + // not be replayed into the replacement merely because it reused the same id. + upstreamRevision: string; kind: UpstreamProviderKind; name: string; disabledPublicModelIds: readonly string[]; diff --git a/packages/test-utils/src/stubs.ts b/packages/test-utils/src/stubs.ts index 6120905b7..94b879212 100644 --- a/packages/test-utils/src/stubs.ts +++ b/packages/test-utils/src/stubs.ts @@ -95,6 +95,7 @@ export const stubModelCandidate = (overrides: { } = {}): ModelCandidate => { const provider = overrides.provider ?? { upstream: 'test-upstream', + upstreamRevision: 'test-revision', kind: 'custom', name: 'Test Upstream', disabledPublicModelIds: [], From 8d74536f9dfcec9169d69bdfb02ecabba57153c8 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:50:36 +0800 Subject: [PATCH 003/157] feat(gateway): extract client-carried affinity at ingress Decrypt recognized carriers before normal gateway processing, retain only target evidence inside the affinity membrane, and rebuild a fresh source payload for each candidate. Restore original opaque state only for an exact target, strip owned state on fallback, and preserve unrecognized values for cascaded gateways. --- .../chat/affinity/chat-completions-ingress.ts | 29 +++++ .../chat/affinity/gemini-ingress.ts | 53 ++++++++ .../data-plane/chat/affinity/ingress_test.ts | 114 ++++++++++++++++++ .../chat/affinity/messages-ingress.ts | 62 ++++++++++ .../src/data-plane/chat/affinity/prepared.ts | 23 ++++ .../chat/affinity/responses-ingress.ts | 77 ++++++++++++ 6 files changed, 358 insertions(+) create mode 100644 packages/gateway/src/data-plane/chat/affinity/chat-completions-ingress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/gemini-ingress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/ingress_test.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/messages-ingress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/prepared.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/responses-ingress.ts diff --git a/packages/gateway/src/data-plane/chat/affinity/chat-completions-ingress.ts b/packages/gateway/src/data-plane/chat/affinity/chat-completions-ingress.ts new file mode 100644 index 000000000..d4b122351 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/chat-completions-ingress.ts @@ -0,0 +1,29 @@ +import type { AffinityCodec } from './codec.ts'; +import { blobForCandidate, ownedAffinities, type PreparedAffinityPayload } from './prepared.ts'; +import type { DecodedAffinityBlob } from './types.ts'; +import type { ChatCompletionsPayload } from '@floway-dev/protocols/chat-completions'; + +export const prepareChatCompletionsAffinity = async ( + payload: ChatCompletionsPayload, + codec: AffinityCodec, +): Promise> => { + const decoded = new Map(); + for (const [index, message] of payload.messages.entries()) { + if (message.role !== 'assistant' || typeof message.reasoning_opaque !== 'string') continue; + decoded.set(index, await codec.unwrap(message.reasoning_opaque)); + } + + return { + affinities: ownedAffinities(decoded.values()), + payloadForCandidate: candidate => { + const candidatePayload = structuredClone(payload); + for (const [index, blob] of decoded) { + const message = candidatePayload.messages[index]; + const selected = blobForCandidate(blob, candidate); + if (selected.present) message.reasoning_opaque = selected.value; + else delete message.reasoning_opaque; + } + return candidatePayload; + }, + }; +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/gemini-ingress.ts b/packages/gateway/src/data-plane/chat/affinity/gemini-ingress.ts new file mode 100644 index 000000000..242b7d589 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/gemini-ingress.ts @@ -0,0 +1,53 @@ +import type { AffinityCodec } from './codec.ts'; +import { blobForCandidate, ownedAffinities, type PreparedAffinityPayload } from './prepared.ts'; +import type { DecodedAffinityBlob } from './types.ts'; +import type { GeminiPart, GeminiPayload } from '@floway-dev/protocols/gemini'; + +interface GeminiBlobLocation { + readonly contentIndex: number; + readonly partIndex: number; + readonly decoded: DecodedAffinityBlob; +} + +export const prepareGeminiAffinity = async ( + payload: GeminiPayload, + codec: AffinityCodec, +): Promise> => { + const locations: GeminiBlobLocation[] = []; + for (const [contentIndex, content] of (payload.contents ?? []).entries()) { + if (content.role !== 'model') continue; + for (const [partIndex, part] of content.parts.entries()) { + if (typeof part.thoughtSignature !== 'string') continue; + locations.push({ contentIndex, partIndex, decoded: await codec.unwrap(part.thoughtSignature) }); + } + } + + return { + affinities: ownedAffinities(locations.map(location => location.decoded)), + payloadForCandidate: candidate => { + const candidatePayload = structuredClone(payload); + if (candidatePayload.contents === undefined) return candidatePayload; + const byContent = Map.groupBy(locations, location => location.contentIndex); + for (const [contentIndex, contentLocations] of byContent) { + const content = candidatePayload.contents[contentIndex]; + const replacements = new Map(); + for (const location of contentLocations) { + const part = content.parts[location.partIndex]; + const selected = blobForCandidate(location.decoded, candidate); + if (selected.present) replacements.set(location.partIndex, { ...part, thoughtSignature: selected.value }); + else if (location.decoded.kind === 'owned' && location.decoded.value === undefined) replacements.set(location.partIndex, null); + else { + const replacement = { ...part }; + delete replacement.thoughtSignature; + replacements.set(location.partIndex, replacement); + } + } + content.parts = content.parts.flatMap((part, partIndex) => { + const replacement = replacements.get(partIndex); + return replacement === undefined ? [part] : replacement === null ? [] : [replacement]; + }); + } + return candidatePayload; + }, + }; +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/ingress_test.ts b/packages/gateway/src/data-plane/chat/affinity/ingress_test.ts new file mode 100644 index 000000000..85f89bcff --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/ingress_test.ts @@ -0,0 +1,114 @@ +import { describe, expect, test } from 'vitest'; + +import { affinityTargetForCandidate } from './candidate.ts'; +import { prepareChatCompletionsAffinity } from './chat-completions-ingress.ts'; +import { AffinityCodec } from './codec.ts'; +import { prepareGeminiAffinity } from './gemini-ingress.ts'; +import { prepareMessagesAffinity } from './messages-ingress.ts'; +import { prepareResponsesAffinity } from './responses-ingress.ts'; +import type { ModelCandidate } from '@floway-dev/provider'; +import { stubModelCandidate } from '@floway-dev/test-utils'; + +const codec = new AffinityCodec('22'.repeat(32)); + +const candidate = (upstream: string): ModelCandidate => { + const base = stubModelCandidate(); + return stubModelCandidate({ + provider: { ...base.provider, upstream, upstreamRevision: `${upstream}-revision` }, + model: { id: 'model' }, + }); +}; + +const candidateA = candidate('upstream-a'); +const candidateB = candidate('upstream-b'); + +describe('protocol affinity ingress', () => { + test('Chat Completions restores owned opaque state only for its exact candidate', async () => { + const carrier = await codec.wrap('upstream-signature', affinityTargetForCandidate(candidateA, 'prefer')); + const prepared = await prepareChatCompletionsAffinity({ + model: 'model', + messages: [{ role: 'assistant', content: 'answer', reasoning_opaque: carrier }], + }, codec); + + expect(prepared.affinities).toEqual([affinityTargetForCandidate(candidateA, 'prefer')]); + expect(prepared.payloadForCandidate(candidateA).messages[0]).toMatchObject({ reasoning_opaque: 'upstream-signature' }); + expect(prepared.payloadForCandidate(candidateB).messages[0]).not.toHaveProperty('reasoning_opaque'); + }); + + test('Messages removes synthetic blocks and strips incompatible signatures without hiding thinking', async () => { + const signature = await codec.wrap('signature', affinityTargetForCandidate(candidateA, 'prefer')); + const synthetic = await codec.wrap(undefined, affinityTargetForCandidate(candidateA, 'prefer')); + const prepared = await prepareMessagesAffinity({ + model: 'model', + max_tokens: 100, + messages: [{ + role: 'assistant', + content: [ + { type: 'thinking', thinking: 'visible reasoning', signature }, + { type: 'redacted_thinking', data: synthetic }, + { type: 'text', text: 'answer' }, + ], + }], + }, codec); + + expect(prepared.payloadForCandidate(candidateA).messages[0]).toEqual({ + role: 'assistant', + content: [ + { type: 'thinking', thinking: 'visible reasoning', signature: 'signature' }, + { type: 'text', text: 'answer' }, + ], + }); + expect(prepared.payloadForCandidate(candidateB).messages[0]).toEqual({ + role: 'assistant', + content: [ + { type: 'thinking', thinking: 'visible reasoning' }, + { type: 'text', text: 'answer' }, + ], + }); + }); + + test('Gemini removes a synthetic signature-only part and preserves foreign signatures', async () => { + const synthetic = await codec.wrap(undefined, affinityTargetForCandidate(candidateA, 'prefer')); + const prepared = await prepareGeminiAffinity({ + contents: [{ + role: 'model', + parts: [ + { text: 'answer' }, + { thoughtSignature: synthetic }, + { text: 'foreign', thoughtSignature: 'not-floway' }, + ], + }], + }, codec); + + expect(prepared.payloadForCandidate(candidateA).contents?.[0].parts).toEqual([ + { text: 'answer' }, + { text: 'foreign', thoughtSignature: 'not-floway' }, + ]); + }); + + test('Responses restores the original upstream item id and drops owned state on fallback', async () => { + const carrier = await codec.wrap('encrypted', affinityTargetForCandidate(candidateA, 'prefer', 'rs_upstream')); + const prepared = await prepareResponsesAffinity({ + model: 'model', + input: [{ type: 'reasoning', id: 'rs_gateway', summary: [{ type: 'summary_text', text: 'visible' }], encrypted_content: carrier }], + }, codec); + + expect(prepared.payloadForCandidate(candidateA).input).toEqual([ + { type: 'reasoning', id: 'rs_upstream', summary: [{ type: 'summary_text', text: 'visible' }], encrypted_content: 'encrypted' }, + ]); + const fallback = prepared.payloadForCandidate(candidateB).input; + expect(fallback).toHaveLength(1); + expect(fallback[0]).not.toHaveProperty('encrypted_content'); + expect((fallback[0] as { id?: string }).id).toMatch(/^rs_tmp_/); + }); + + test('foreign blobs pass through unchanged for cascaded gateways', async () => { + const prepared = await prepareResponsesAffinity({ + model: 'model', + input: [{ type: 'reasoning', id: 'rs_foreign', summary: [], encrypted_content: 'foreign' }], + }, codec); + + expect(prepared.affinities).toEqual([]); + expect(prepared.payloadForCandidate(candidateA).input[0]).toMatchObject({ encrypted_content: 'foreign' }); + }); +}); diff --git a/packages/gateway/src/data-plane/chat/affinity/messages-ingress.ts b/packages/gateway/src/data-plane/chat/affinity/messages-ingress.ts new file mode 100644 index 000000000..c7fadd2aa --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/messages-ingress.ts @@ -0,0 +1,62 @@ +import type { AffinityCodec } from './codec.ts'; +import { blobForCandidate, ownedAffinities, type PreparedAffinityPayload } from './prepared.ts'; +import type { DecodedAffinityBlob } from './types.ts'; +import type { MessagesAssistantContentBlock, MessagesPayload } from '@floway-dev/protocols/messages'; + +interface MessagesBlobLocation { + readonly messageIndex: number; + readonly blockIndex: number; + readonly kind: 'thinking' | 'redacted_thinking'; + readonly decoded: DecodedAffinityBlob; +} + +export const prepareMessagesAffinity = async ( + payload: MessagesPayload, + codec: AffinityCodec, +): Promise> => { + const locations: MessagesBlobLocation[] = []; + for (const [messageIndex, message] of payload.messages.entries()) { + if (message.role !== 'assistant' || !Array.isArray(message.content)) continue; + for (const [blockIndex, block] of message.content.entries()) { + if (block.type === 'thinking' && typeof block.signature === 'string') { + locations.push({ messageIndex, blockIndex, kind: block.type, decoded: await codec.unwrap(block.signature) }); + } else if (block.type === 'redacted_thinking') { + locations.push({ messageIndex, blockIndex, kind: block.type, decoded: await codec.unwrap(block.data) }); + } + } + } + + return { + affinities: ownedAffinities(locations.map(location => location.decoded)), + payloadForCandidate: candidate => { + const candidatePayload = structuredClone(payload); + const byMessage = Map.groupBy(locations, location => location.messageIndex); + for (const [messageIndex, messageLocations] of byMessage) { + const message = candidatePayload.messages[messageIndex]; + if (message.role !== 'assistant' || !Array.isArray(message.content)) { + throw new Error('Messages affinity location no longer points to assistant block content'); + } + const replacements = new Map(); + for (const location of messageLocations) { + const block = message.content[location.blockIndex]; + const selected = blobForCandidate(location.decoded, candidate); + if (location.kind === 'thinking') { + if (block.type !== 'thinking') throw new Error('Messages affinity thinking location changed type'); + const replacement = { ...block }; + if (selected.present) replacement.signature = selected.value; + else delete replacement.signature; + replacements.set(location.blockIndex, replacement); + } else { + if (block.type !== 'redacted_thinking') throw new Error('Messages affinity redacted location changed type'); + replacements.set(location.blockIndex, selected.present ? { ...block, data: selected.value } : null); + } + } + message.content = message.content.flatMap((block, blockIndex) => { + const replacement = replacements.get(blockIndex); + return replacement === undefined ? [block] : replacement === null ? [] : [replacement]; + }); + } + return candidatePayload; + }, + }; +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/prepared.ts b/packages/gateway/src/data-plane/chat/affinity/prepared.ts new file mode 100644 index 000000000..2d3c6567a --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/prepared.ts @@ -0,0 +1,23 @@ +import { candidateMatchesAffinity } from './candidate.ts'; +import type { AffinityTarget, DecodedAffinityBlob } from './types.ts'; +import type { ModelCandidate } from '@floway-dev/provider'; + +export interface PreparedAffinityPayload { + readonly affinities: readonly AffinityTarget[]; + readonly payloadForCandidate: (candidate: ModelCandidate) => T; +} + +export type CandidateBlob = + | { readonly present: false; readonly owned: true } + | { readonly present: true; readonly owned: boolean; readonly value: string }; + +export const blobForCandidate = (decoded: DecodedAffinityBlob, candidate: ModelCandidate): CandidateBlob => { + if (decoded.kind === 'foreign') return { present: true, owned: false, value: decoded.value }; + if (!candidateMatchesAffinity(candidate, decoded.envelope.affinity) || decoded.value === undefined) { + return { present: false, owned: true }; + } + return { present: true, owned: true, value: decoded.value }; +}; + +export const ownedAffinities = (decoded: Iterable): AffinityTarget[] => + [...decoded].flatMap(blob => blob.kind === 'owned' ? [blob.envelope.affinity] : []); diff --git a/packages/gateway/src/data-plane/chat/affinity/responses-ingress.ts b/packages/gateway/src/data-plane/chat/affinity/responses-ingress.ts new file mode 100644 index 000000000..e37d2baa4 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/responses-ingress.ts @@ -0,0 +1,77 @@ +import type { AffinityCodec } from './codec.ts'; +import { blobForCandidate, ownedAffinities, type PreparedAffinityPayload } from './prepared.ts'; +import type { DecodedAffinityBlob } from './types.ts'; +import { createTemporaryResponsesItemId } from '../responses/items/format.ts'; +import type { CanonicalResponsesPayload, ResponsesInputItem } from '@floway-dev/protocols/responses'; + +interface ResponsesBlobLocation { + readonly itemIndex: number; + readonly contentIndex?: number; + readonly decoded: DecodedAffinityBlob; +} + +const encryptedContentLocations = async ( + items: readonly ResponsesInputItem[], + codec: AffinityCodec, +): Promise => { + const locations: ResponsesBlobLocation[] = []; + for (const [itemIndex, item] of items.entries()) { + const topLevel = (item as { encrypted_content?: unknown }).encrypted_content; + if (typeof topLevel === 'string') locations.push({ itemIndex, decoded: await codec.unwrap(topLevel) }); + if (item.type !== 'agent_message') continue; + for (const [contentIndex, content] of item.content.entries()) { + if (content.type !== 'encrypted_content' || typeof content.encrypted_content !== 'string') continue; + locations.push({ itemIndex, contentIndex, decoded: await codec.unwrap(content.encrypted_content) }); + } + } + return locations; +}; + +export const prepareResponsesAffinity = async ( + payload: CanonicalResponsesPayload, + codec: AffinityCodec, +): Promise> => { + const locations = await encryptedContentLocations(payload.input, codec); + return { + affinities: ownedAffinities(locations.map(location => location.decoded)), + payloadForCandidate: candidate => { + const candidatePayload = structuredClone(payload); + const byItem = Map.groupBy(locations, location => location.itemIndex); + const rewritten = candidatePayload.input.flatMap((item, itemIndex): ResponsesInputItem[] => { + const itemLocations = byItem.get(itemIndex); + if (itemLocations === undefined) return [item]; + let removeItem = false; + const replacement = { ...item } as ResponsesInputItem & { encrypted_content?: string }; + for (const location of itemLocations) { + const selected = blobForCandidate(location.decoded, candidate); + if (location.contentIndex === undefined) { + if (selected.present) { + replacement.encrypted_content = selected.value; + if (selected.owned && location.decoded.kind === 'owned') { + const upstreamItemId = location.decoded.envelope.affinity.upstreamItemId; + if (upstreamItemId !== undefined) replacement.id = upstreamItemId; + } + } else { + delete replacement.encrypted_content; + if (location.decoded.kind === 'owned' && location.decoded.value === undefined) removeItem = true; + else if ('id' in replacement && typeof replacement.id === 'string') { + replacement.id = createTemporaryResponsesItemId(replacement.type); + } + } + continue; + } + if (replacement.type !== 'agent_message') throw new Error('Responses affinity content location changed item type'); + replacement.content = replacement.content.flatMap((content, contentIndex) => { + if (contentIndex !== location.contentIndex) return [content]; + if (content.type !== 'encrypted_content') throw new Error('Responses affinity content location changed block type'); + return selected.present ? [{ ...content, encrypted_content: selected.value }] : []; + }); + } + if (removeItem) return []; + if (replacement.type === 'agent_message' && replacement.content.length === 0) return []; + return [replacement]; + }); + return { ...candidatePayload, input: rewritten }; + }, + }; +}; From b2ec1861fdc565078ed47a105401c5f8e26a87df Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:52:07 +0800 Subject: [PATCH 004/157] feat(gateway): route chat sources by client affinity Build a request-scoped codec from the authenticated key, decrypt source-protocol carriers before candidate enumeration, and supply each fallback attempt with a fresh candidate-specific payload. Exact force/prefer routing now applies to Chat Completions, Messages, Gemini, and their count endpoints without exposing the raw secret to providers. --- .../src/data-plane/chat/affinity/context.ts | 22 +++++++++++++++ .../data-plane/chat/chat-completions/serve.ts | 18 ++++++------ .../src/data-plane/chat/gemini/serve.ts | 28 ++++++++----------- .../src/data-plane/chat/messages/serve.ts | 28 ++++++++----------- .../src/data-plane/chat/shared/gateway-ctx.ts | 8 +++++- .../gateway/src/test-helpers/gateway-ctx.ts | 2 ++ 6 files changed, 64 insertions(+), 42 deletions(-) create mode 100644 packages/gateway/src/data-plane/chat/affinity/context.ts diff --git a/packages/gateway/src/data-plane/chat/affinity/context.ts b/packages/gateway/src/data-plane/chat/affinity/context.ts new file mode 100644 index 000000000..3c180529b --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/context.ts @@ -0,0 +1,22 @@ +import { affinityTargetForCandidate } from './candidate.ts'; +import { AffinityCodec } from './codec.ts'; +import type { AffinityTarget } from './types.ts'; +import type { ModelCandidate } from '@floway-dev/provider'; + +export class AffinityRequestContext { + readonly codec: AffinityCodec; + #selectedCandidate: ModelCandidate | undefined; + + constructor(secret: string) { + this.codec = new AffinityCodec(secret); + } + + select(candidate: ModelCandidate): void { + this.#selectedCandidate = candidate; + } + + selectedTarget(mode: AffinityTarget['mode'] = 'prefer'): AffinityTarget { + if (this.#selectedCandidate === undefined) throw new Error('Affinity target requested before a candidate was selected'); + return affinityTargetForCandidate(this.#selectedCandidate, mode); + } +} diff --git a/packages/gateway/src/data-plane/chat/chat-completions/serve.ts b/packages/gateway/src/data-plane/chat/chat-completions/serve.ts index c2f7b2ba7..8913d1a04 100644 --- a/packages/gateway/src/data-plane/chat/chat-completions/serve.ts +++ b/packages/gateway/src/data-plane/chat/chat-completions/serve.ts @@ -1,14 +1,14 @@ import { chatCompletionsAttempt, chatCompletionsTarget } from './attempt.ts'; import { renderChatCompletionsFailure } from './errors.ts'; +import { prepareChatCompletionsAffinity } from '../affinity/chat-completions-ingress.ts'; +import { routeCandidatesByAffinity } from '../affinity/candidate.ts'; import { enumerateModelCandidates } from '../../providers/registry.ts'; import { iterateCandidates } from '../../shared/iterate-candidates.ts'; -import { classifyResponsesItemAffinity } from '../responses/items/affinity.ts'; import { noViableCandidateFailure } from '../shared/errors.ts'; import type { ChatGatewayCtx } from '../shared/gateway-ctx.ts'; import type { ChatCompletionsPayload, ChatCompletionsStreamEvent } from '@floway-dev/protocols/chat-completions'; import type { ProtocolFrame } from '@floway-dev/protocols/common'; import type { ExecuteResult } from '@floway-dev/provider'; -import { chatCompletionsViaResponsesItemsView } from '@floway-dev/translate/via-responses/responses-items'; export interface ChatCompletionsServeGenerateArgs { readonly payload: ChatCompletionsPayload; @@ -19,6 +19,7 @@ export interface ChatCompletionsServeGenerateArgs { export const chatCompletionsServe = { generate: async (args: ChatCompletionsServeGenerateArgs): Promise>> => { const { payload, ctx, headers } = args; + const prepared = await prepareChatCompletionsAffinity(payload, ctx.affinity.codec); const { candidates: enumerated, sawModel, failedUpstreams } = await enumerateModelCandidates({ upstreamIds: ctx.upstreamIds, model: payload.model, @@ -27,12 +28,7 @@ export const chatCompletionsServe = { runtimeLocation: ctx.runtimeLocation, }); const viable = enumerated.filter(c => chatCompletionsTarget.canServe(c.model.endpoints)); - const decision = await classifyResponsesItemAffinity({ - sourceItems: payload.messages, - view: chatCompletionsViaResponsesItemsView, - store: ctx.store, - candidates: viable, - }); + const decision = routeCandidatesByAffinity(viable, prepared.affinities); if (decision.kind === 'failure') return renderChatCompletionsFailure(decision.failure); if (decision.candidates.length === 0) return renderChatCompletionsFailure(noViableCandidateFailure(sawModel, payload.model, failedUpstreams)); @@ -49,7 +45,11 @@ export const chatCompletionsServe = { 'chatCompletionsServe.generate', ctx, 'chat', - candidate => chatCompletionsAttempt.generate({ payload, ctx, candidate, headers }), + async candidate => { + const result = await chatCompletionsAttempt.generate({ payload: prepared.payloadForCandidate(candidate), ctx, candidate, headers }); + if (result.type === 'events') ctx.affinity.select(candidate); + return result; + }, ); }, }; diff --git a/packages/gateway/src/data-plane/chat/gemini/serve.ts b/packages/gateway/src/data-plane/chat/gemini/serve.ts index deecf1750..d2d7685d3 100644 --- a/packages/gateway/src/data-plane/chat/gemini/serve.ts +++ b/packages/gateway/src/data-plane/chat/gemini/serve.ts @@ -1,14 +1,14 @@ import { geminiAttempt, geminiCountTokensTarget, geminiGenerateTarget } from './attempt.ts'; import { renderGeminiFailure } from './errors.ts'; +import { routeCandidatesByAffinity } from '../affinity/candidate.ts'; +import { prepareGeminiAffinity } from '../affinity/gemini-ingress.ts'; import { enumerateModelCandidates } from '../../providers/registry.ts'; import { iterateCandidates } from '../../shared/iterate-candidates.ts'; -import { classifyResponsesItemAffinity } from '../responses/items/affinity.ts'; import { noViableCandidateFailure } from '../shared/errors.ts'; import type { ChatGatewayCtx } from '../shared/gateway-ctx.ts'; import type { ProtocolFrame } from '@floway-dev/protocols/common'; import type { GeminiPayload, GeminiStreamEvent } from '@floway-dev/protocols/gemini'; import type { ExecuteResult, PlainResult } from '@floway-dev/provider'; -import { geminiViaResponsesItemsView } from '@floway-dev/translate/via-responses/responses-items'; export interface GeminiServeGenerateArgs { readonly payload: GeminiPayload; @@ -30,6 +30,7 @@ export interface GeminiServeCountTokensArgs { export const geminiServe = { generate: async (args: GeminiServeGenerateArgs): Promise>> => { const { payload, ctx, model, headers } = args; + const prepared = await prepareGeminiAffinity(payload, ctx.affinity.codec); const { candidates: enumerated, sawModel, failedUpstreams } = await enumerateModelCandidates({ upstreamIds: ctx.upstreamIds, model, @@ -38,12 +39,7 @@ export const geminiServe = { runtimeLocation: ctx.runtimeLocation, }); const viable = enumerated.filter(c => geminiGenerateTarget.canServe(c.model.endpoints)); - const decision = await classifyResponsesItemAffinity({ - sourceItems: payload.contents ?? [], - view: geminiViaResponsesItemsView, - store: ctx.store, - candidates: viable, - }); + const decision = routeCandidatesByAffinity(viable, prepared.affinities); if (decision.kind === 'failure') return renderGeminiFailure(decision.failure, 'generate'); if (decision.candidates.length === 0) return renderGeminiFailure(noViableCandidateFailure(sawModel, model, failedUpstreams), 'generate'); @@ -61,12 +57,17 @@ export const geminiServe = { 'geminiServe.generate', ctx, 'chat', - candidate => geminiAttempt.generate({ payload, ctx, candidate, headers }), + async candidate => { + const result = await geminiAttempt.generate({ payload: prepared.payloadForCandidate(candidate), ctx, candidate, headers }); + if (result.type === 'events') ctx.affinity.select(candidate); + return result; + }, ); }, countTokens: async (args: GeminiServeCountTokensArgs): Promise> | PlainResult> => { const { payload, ctx, model, headers } = args; + const prepared = await prepareGeminiAffinity(payload, ctx.affinity.codec); const { candidates: enumerated, sawModel, failedUpstreams } = await enumerateModelCandidates({ upstreamIds: ctx.upstreamIds, model, @@ -75,12 +76,7 @@ export const geminiServe = { runtimeLocation: ctx.runtimeLocation, }); const viable = enumerated.filter(c => geminiCountTokensTarget.canServe(c.model.endpoints)); - const decision = await classifyResponsesItemAffinity({ - sourceItems: payload.contents ?? [], - view: geminiViaResponsesItemsView, - store: ctx.store, - candidates: viable, - }); + const decision = routeCandidatesByAffinity(viable, prepared.affinities); if (decision.kind === 'failure') return renderGeminiFailure(decision.failure, 'countTokens'); if (decision.candidates.length === 0) return renderGeminiFailure(noViableCandidateFailure(sawModel, model, failedUpstreams), 'countTokens'); @@ -89,7 +85,7 @@ export const geminiServe = { 'geminiServe.countTokens', ctx, 'chat', - candidate => geminiAttempt.countTokens({ payload, ctx, candidate, headers }), + candidate => geminiAttempt.countTokens({ payload: prepared.payloadForCandidate(candidate), ctx, candidate, headers }), ); }, }; diff --git a/packages/gateway/src/data-plane/chat/messages/serve.ts b/packages/gateway/src/data-plane/chat/messages/serve.ts index 3fd085569..0143dc2eb 100644 --- a/packages/gateway/src/data-plane/chat/messages/serve.ts +++ b/packages/gateway/src/data-plane/chat/messages/serve.ts @@ -1,14 +1,14 @@ import { messagesAttempt, messagesGenerateTarget, messagesCountTokensTarget } from './attempt.ts'; import { renderMessagesFailure } from './errors.ts'; +import { routeCandidatesByAffinity } from '../affinity/candidate.ts'; +import { prepareMessagesAffinity } from '../affinity/messages-ingress.ts'; import { enumerateModelCandidates } from '../../providers/registry.ts'; import { iterateCandidates } from '../../shared/iterate-candidates.ts'; -import { classifyResponsesItemAffinity } from '../responses/items/affinity.ts'; import { noViableCandidateFailure } from '../shared/errors.ts'; import type { ChatGatewayCtx } from '../shared/gateway-ctx.ts'; import type { ProtocolFrame } from '@floway-dev/protocols/common'; import type { MessagesPayload, MessagesStreamEvent } from '@floway-dev/protocols/messages'; import type { ExecuteResult, PlainResult } from '@floway-dev/provider'; -import { messagesViaResponsesItemsView } from '@floway-dev/translate/via-responses/responses-items'; export interface MessagesServeGenerateArgs { readonly payload: MessagesPayload; @@ -25,6 +25,7 @@ export interface MessagesServeCountTokensArgs { export const messagesServe = { generate: async (args: MessagesServeGenerateArgs): Promise>> => { const { payload, ctx, headers } = args; + const prepared = await prepareMessagesAffinity(payload, ctx.affinity.codec); const { candidates: enumerated, sawModel, failedUpstreams } = await enumerateModelCandidates({ upstreamIds: ctx.upstreamIds, model: payload.model, @@ -33,12 +34,7 @@ export const messagesServe = { runtimeLocation: ctx.runtimeLocation, }); const viable = enumerated.filter(c => messagesGenerateTarget.canServe(c.model.endpoints)); - const decision = await classifyResponsesItemAffinity({ - sourceItems: payload.messages, - view: messagesViaResponsesItemsView, - store: ctx.store, - candidates: viable, - }); + const decision = routeCandidatesByAffinity(viable, prepared.affinities); if (decision.kind === 'failure') return renderMessagesFailure(decision.failure, 'generate'); if (decision.candidates.length === 0) return renderMessagesFailure(noViableCandidateFailure(sawModel, payload.model, failedUpstreams), 'generate'); @@ -53,12 +49,17 @@ export const messagesServe = { 'messagesServe.generate', ctx, 'chat', - candidate => messagesAttempt.generate({ payload, ctx, candidate, headers }), + async candidate => { + const result = await messagesAttempt.generate({ payload: prepared.payloadForCandidate(candidate), ctx, candidate, headers }); + if (result.type === 'events') ctx.affinity.select(candidate); + return result; + }, ); }, countTokens: async (args: MessagesServeCountTokensArgs): Promise> | PlainResult> => { const { payload, ctx, headers } = args; + const prepared = await prepareMessagesAffinity(payload, ctx.affinity.codec); const { candidates: enumerated, sawModel, failedUpstreams } = await enumerateModelCandidates({ upstreamIds: ctx.upstreamIds, model: payload.model, @@ -67,12 +68,7 @@ export const messagesServe = { runtimeLocation: ctx.runtimeLocation, }); const viable = enumerated.filter(c => messagesCountTokensTarget.canServe(c.model.endpoints)); - const decision = await classifyResponsesItemAffinity({ - sourceItems: payload.messages, - view: messagesViaResponsesItemsView, - store: ctx.store, - candidates: viable, - }); + const decision = routeCandidatesByAffinity(viable, prepared.affinities); if (decision.kind === 'failure') return renderMessagesFailure(decision.failure, 'countTokens'); if (decision.candidates.length === 0) return renderMessagesFailure(noViableCandidateFailure(sawModel, payload.model, failedUpstreams), 'countTokens'); @@ -81,7 +77,7 @@ export const messagesServe = { 'messagesServe.countTokens', ctx, 'chat', - candidate => messagesAttempt.countTokens({ payload, ctx, candidate, headers }), + candidate => messagesAttempt.countTokens({ payload: prepared.payloadForCandidate(candidate), ctx, candidate, headers }), ); }, }; diff --git a/packages/gateway/src/data-plane/chat/shared/gateway-ctx.ts b/packages/gateway/src/data-plane/chat/shared/gateway-ctx.ts index 7e9f29349..50777b63a 100644 --- a/packages/gateway/src/data-plane/chat/shared/gateway-ctx.ts +++ b/packages/gateway/src/data-plane/chat/shared/gateway-ctx.ts @@ -1,4 +1,5 @@ import type { RequestBody } from './request-body.ts'; +import { AffinityRequestContext } from '../affinity/context.ts'; import { type DumpAccumulator, openDumpAccumulator } from '../../../dump/accumulator.ts'; import { apiKeyFromContext, type AuthedContext, effectiveUpstreamIdsFromContext } from '../../../middleware/auth.ts'; import { getRuntimeLocation } from '../../../runtime/runtime-info.ts'; @@ -54,6 +55,7 @@ export interface GatewayCtx { // attempt. Passthrough endpoints (embeddings / images / completions) have // no stored-items concept and stay on plain `GatewayCtx`. export interface ChatGatewayCtx extends GatewayCtx { + readonly affinity: AffinityRequestContext; readonly store: StatefulResponsesStore; } @@ -131,5 +133,9 @@ export const createChatGatewayCtxFromHono = ( storeFactory: (apiKeyId: string) => StatefulResponsesStore, ): ChatGatewayCtx => { const base = createGatewayCtxFromHono(c, opts); - return { ...base, store: storeFactory(base.apiKeyId) }; + return { + ...base, + affinity: new AffinityRequestContext(apiKeyFromContext(c).affinitySecret), + store: storeFactory(base.apiKeyId), + }; }; diff --git a/packages/gateway/src/test-helpers/gateway-ctx.ts b/packages/gateway/src/test-helpers/gateway-ctx.ts index 5c6d995b2..be20dd513 100644 --- a/packages/gateway/src/test-helpers/gateway-ctx.ts +++ b/packages/gateway/src/test-helpers/gateway-ctx.ts @@ -1,4 +1,5 @@ import { createNonResponsesSourceStore } from '../data-plane/chat/responses/items/store.ts'; +import { AffinityRequestContext } from '../data-plane/chat/affinity/context.ts'; import type { ChatGatewayCtx, GatewayCtx } from '../data-plane/chat/shared/gateway-ctx.ts'; // Shared minimal GatewayCtx for tests that exercise serve / respond / @@ -27,6 +28,7 @@ export const mockChatGatewayCtx = (overrides: Partial = {}): Cha const base = mockGatewayCtx(overrides); return { ...base, + affinity: overrides.affinity ?? new AffinityRequestContext('00'.repeat(32)), store: overrides.store ?? createNonResponsesSourceStore(base.apiKeyId), }; }; From 8f1cf611366fd67a18d4bfa539ae8ca5f11d6476 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:52:10 +0800 Subject: [PATCH 005/157] fix(protocols): treat opaque reasoning as snapshots Align Anthropic signature_delta and Chat reasoning_opaque with last-write-wins snapshot semantics. Reassemble every Chat choice independently, including per-choice extras, tools, reasoning, and finish state, while retaining response-level usage and metadata. Update Chat-target translators so repeated source signatures replace rather than concatenate. --- .../src/chat-completions/reassemble.ts | 229 ++++++++---------- .../src/chat-completions/reassemble_test.ts | 105 +++++++- packages/protocols/src/messages/reassemble.ts | 2 +- .../protocols/src/messages/reassemble_test.ts | 8 + .../src/gemini-via-chat-completions/events.ts | 4 +- .../events_test.ts | 1 + .../gemini-via-chat-completions/request.ts | 4 +- .../request_test.ts | 2 +- .../src/gemini-via-messages/events.ts | 6 +- .../src/gemini-via-messages/events_test.ts | 5 + .../messages-via-chat-completions/events.ts | 6 +- .../events_test.ts | 3 +- .../translate/src/shared/gemini-via/gemini.ts | 4 +- 13 files changed, 233 insertions(+), 146 deletions(-) diff --git a/packages/protocols/src/chat-completions/reassemble.ts b/packages/protocols/src/chat-completions/reassemble.ts index d4bfe2267..000ff8200 100644 --- a/packages/protocols/src/chat-completions/reassemble.ts +++ b/packages/protocols/src/chat-completions/reassemble.ts @@ -1,168 +1,143 @@ import { chatCompletionsErrorPayloadMessage } from './index.ts'; -import type { ChatCompletionsStreamEvent, ChatCompletionsResult, ChatCompletionsReasoningItem, ChatCompletionsToolCall } from './index.ts'; +import type { ChatCompletionsChoiceNonStreaming, ChatCompletionsResult, ChatCompletionsStreamEvent, ChatCompletionsReasoningItem, ChatCompletionsToolCall } from './index.ts'; import { captureExtras } from '../common/reassemble-extras.ts'; -// Field-fidelity contract: every field a Chat Completions upstream emits must -// reach the non-streaming client untouched. The stream path achieves this for -// free (`to-sse.ts` re-serialises each chunk via `JSON.stringify`), but -// non-streaming clients receive a single result object reassembled here, so -// any field this code does not explicitly carry over disappears. -// -// We split fields into two buckets: -// - "known" fields with streaming semantics (string concat, array merge by -// index) — handled by the typed accumulators below. -// - everything else — captured generically via {@link captureExtras} and -// replayed onto the assembled result. -// -// `reasoning_content` (the DeepSeek/Kimi dialect for reasoning text on Chat -// Completions) is the concrete case that motivated the extras path; the goal -// is that any future field — `audio_prompt_tokens`, `prompt_filter_results`, -// `this_is_a_non_standard_field_of_reasoning` — survives by default without a -// gateway code change. -// -// Stable scalar strings that the upstream repeats unchanged on every chunk — -// OpenAI's `system_fingerprint` and `service_tier` are the canonical examples -// — MUST be registered as known keys; otherwise the generic extras path -// concatenates them into a duplicated mess (`fp_xfp_xfp_x…`). - +// Field-fidelity contract: every field an upstream emits must reach the +// non-streaming result. Known streaming fields use their protocol semantics; +// unknown fields fall through to captureExtras so future extensions survive. const KNOWN_DELTA_KEYS = new Set(['content', 'role', 'reasoning_text', 'reasoning_opaque', 'reasoning_items', 'tool_calls']); const KNOWN_CHOICE_KEYS = new Set(['index', 'delta', 'finish_reason']); const KNOWN_CHUNK_KEYS = new Set(['id', 'object', 'created', 'model', 'choices', 'usage', 'system_fingerprint', 'service_tier']); +interface ToolCallAccumulator { + id: string; + name: string; + arguments: string; +} + +interface ChoiceAccumulator { + readonly index: number; + content: string; + reasoningText: string; + reasoningOpaque?: string; + readonly reasoningItems: ChatCompletionsReasoningItem[]; + finishReason: ChatCompletionsChoiceNonStreaming['finish_reason']; + readonly toolCalls: Map; + readonly choiceExtras: Record; + readonly messageExtras: Record; +} + +const createChoiceAccumulator = (index: number): ChoiceAccumulator => ({ + index, + content: '', + reasoningText: '', + reasoningItems: [], + finishReason: 'stop', + toolCalls: new Map(), + choiceExtras: {}, + messageExtras: {}, +}); + +const accumulateToolCalls = (choice: ChoiceAccumulator, value: unknown): void => { + if (!Array.isArray(value)) return; + + for (const raw of value) { + if (raw === null || typeof raw !== 'object') continue; + const toolCall = raw as Record; + if (typeof toolCall.index !== 'number') continue; + const fn = toolCall.function && typeof toolCall.function === 'object' + ? toolCall.function as Record + : undefined; + const current = choice.toolCalls.get(toolCall.index) ?? { id: '', name: '', arguments: '' }; + if (typeof toolCall.id === 'string') current.id = toolCall.id; + if (typeof fn?.name === 'string') current.name = fn.name; + if (typeof fn?.arguments === 'string') current.arguments += fn.arguments; + choice.toolCalls.set(toolCall.index, current); + } +}; + +const finalizedToolCalls = (choice: ChoiceAccumulator): ChatCompletionsToolCall[] => + [...choice.toolCalls.entries()] + .toSorted(([left], [right]) => left - right) + .map(([, toolCall]) => ({ + id: toolCall.id, + type: 'function', + function: { name: toolCall.name, arguments: toolCall.arguments }, + })); + +const finalizeChoice = (choice: ChoiceAccumulator): ChatCompletionsChoiceNonStreaming => { + const toolCalls = finalizedToolCalls(choice); + return { + index: choice.index, + message: { + role: 'assistant', + content: choice.content || null, + ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}), + ...(choice.reasoningText ? { reasoning_text: choice.reasoningText } : {}), + ...(choice.reasoningOpaque !== undefined ? { reasoning_opaque: choice.reasoningOpaque } : {}), + ...(choice.reasoningItems.length > 0 ? { reasoning_items: choice.reasoningItems } : {}), + ...choice.messageExtras, + }, + finish_reason: choice.finishReason, + ...choice.choiceExtras, + } as ChatCompletionsChoiceNonStreaming; +}; + export async function reassembleChatCompletionsEvents(chunks: AsyncIterable): Promise { let id = ''; let model = ''; let created = 0; let systemFingerprint: string | undefined; let serviceTier: ChatCompletionsResult['service_tier']; - let content = ''; - let reasoningText = ''; - let reasoningOpaque = ''; - let hasReasoningOpaque = false; - const reasoningItems: ChatCompletionsReasoningItem[] = []; - let finishReason: ChatCompletionsResult['choices'][number]['finish_reason'] = 'stop'; let lastUsage: ChatCompletionsResult['usage'] | undefined; - - const toolCallsMap = new Map(); - + const choices = new Map(); const chunkExtras: Record = {}; - const choiceExtras: Record = {}; - const messageExtras: Record = {}; for await (const chunk of chunks) { const errorMessage = chatCompletionsErrorPayloadMessage(chunk); - if (errorMessage) { - throw new Error(`Upstream Chat Completions SSE error: ${errorMessage}`); - } + if (errorMessage) throw new Error(`Upstream Chat Completions SSE error: ${errorMessage}`); if (!id && chunk.id) { - id = chunk.id as string; - model = chunk.model as string; - created = chunk.created as number; + id = chunk.id; + model = chunk.model; + created = chunk.created; } - if (!systemFingerprint && typeof chunk.system_fingerprint === 'string' && chunk.system_fingerprint) { systemFingerprint = chunk.system_fingerprint; } if (!serviceTier && typeof chunk.service_tier === 'string' && chunk.service_tier) { serviceTier = chunk.service_tier; } - - if (chunk.usage) { - lastUsage = chunk.usage as ChatCompletionsResult['usage']; - } - + if (chunk.usage) lastUsage = chunk.usage; captureExtras(chunk as unknown as Record, KNOWN_CHUNK_KEYS, chunkExtras); - const choices = chunk.choices as unknown as Array> | undefined; - if (!choices) continue; - - for (const choice of choices) { - captureExtras(choice, KNOWN_CHOICE_KEYS, choiceExtras); - - const delta = choice.delta as Record | undefined; - if (delta) { - captureExtras(delta, KNOWN_DELTA_KEYS, messageExtras); - - if (typeof delta.content === 'string') { - content += delta.content; - } - if (typeof delta.reasoning_text === 'string') { - reasoningText += delta.reasoning_text; - } - if (typeof delta.reasoning_opaque === 'string') { - reasoningOpaque += delta.reasoning_opaque; - hasReasoningOpaque = true; - } - if (Array.isArray(delta.reasoning_items)) { - reasoningItems.push(...(delta.reasoning_items as ChatCompletionsReasoningItem[])); - } - - if (Array.isArray(delta.tool_calls)) { - for (const toolCall of delta.tool_calls as Array>) { - const idx = toolCall.index as number; - const existing = toolCallsMap.get(idx); - if (!existing) { - toolCallsMap.set(idx, { - id: (toolCall.id as string) ?? '', - name: ((toolCall.function as Record)?.name as string) ?? '', - arguments: ((toolCall.function as Record)?.arguments as string) ?? '', - }); - } else { - if (toolCall.id) existing.id = toolCall.id as string; - const fn = toolCall.function as Record | undefined; - if (fn?.name) existing.name = fn.name as string; - if (fn?.arguments) { - existing.arguments += fn.arguments as string; - } - } - } - } - } - - if (choice.finish_reason) { - finishReason = choice.finish_reason as ChatCompletionsResult['choices'][number]['finish_reason']; + for (const streamed of chunk.choices ?? []) { + const choice = choices.get(streamed.index) ?? createChoiceAccumulator(streamed.index); + choices.set(streamed.index, choice); + captureExtras(streamed as unknown as Record, KNOWN_CHOICE_KEYS, choice.choiceExtras); + + const delta = streamed.delta as unknown as Record; + captureExtras(delta, KNOWN_DELTA_KEYS, choice.messageExtras); + if (typeof delta.content === 'string') choice.content += delta.content; + if (typeof delta.reasoning_text === 'string') choice.reasoningText += delta.reasoning_text; + if (typeof delta.reasoning_opaque === 'string') choice.reasoningOpaque = delta.reasoning_opaque; + if (Array.isArray(delta.reasoning_items)) { + choice.reasoningItems.push(...delta.reasoning_items as ChatCompletionsReasoningItem[]); } + accumulateToolCalls(choice, delta.tool_calls); + if (streamed.finish_reason !== null) choice.finishReason = streamed.finish_reason; } } - const toolCalls: ChatCompletionsToolCall[] = []; - const sortedIndices = [...toolCallsMap.keys()].sort((a, b) => a - b); - for (const idx of sortedIndices) { - const toolCall = toolCallsMap.get(idx)!; - toolCalls.push({ - id: toolCall.id, - type: 'function', - function: { name: toolCall.name, arguments: toolCall.arguments }, - }); - } - - const message = { - role: 'assistant' as const, - content: content || null, - ...(toolCalls.length > 0 && { tool_calls: toolCalls }), - ...(reasoningText && { reasoning_text: reasoningText }), - ...(hasReasoningOpaque ? { reasoning_opaque: reasoningOpaque } : {}), - ...(reasoningItems.length > 0 && { reasoning_items: reasoningItems }), - ...messageExtras, - }; - return { id, object: 'chat.completion', created, model, - choices: [ - { - index: 0, - message, - finish_reason: finishReason, - ...choiceExtras, - }, - ], - ...(systemFingerprint && { system_fingerprint: systemFingerprint }), - ...(serviceTier && { service_tier: serviceTier }), - ...(lastUsage && { usage: lastUsage }), + choices: [...choices.values()].toSorted((left, right) => left.index - right.index).map(finalizeChoice), + ...(systemFingerprint ? { system_fingerprint: systemFingerprint } : {}), + ...(serviceTier ? { service_tier: serviceTier } : {}), + ...(lastUsage ? { usage: lastUsage } : {}), ...chunkExtras, } as ChatCompletionsResult; } diff --git a/packages/protocols/src/chat-completions/reassemble_test.ts b/packages/protocols/src/chat-completions/reassemble_test.ts index 63ab06386..9799fea6a 100644 --- a/packages/protocols/src/chat-completions/reassemble_test.ts +++ b/packages/protocols/src/chat-completions/reassemble_test.ts @@ -141,7 +141,7 @@ test('reassembleChatCompletionsEvents reassembles tool calls', async () => { assertEquals(result.choices[0].message.tool_calls![0].function.arguments, '{"city":"Tokyo"}'); }); -test('reassembleChatCompletionsEvents reassembles reasoning fields', async () => { +test('reassembleChatCompletionsEvents concatenates reasoning text and keeps the latest opaque snapshot', async () => { const body = makeEvents([ { data: { @@ -155,7 +155,7 @@ test('reassembleChatCompletionsEvents reassembles reasoning fields', async () => delta: { role: 'assistant', reasoning_text: 'think', - reasoning_opaque: 'enc', + reasoning_opaque: 'enc_old', }, finish_reason: null, }, @@ -171,7 +171,7 @@ test('reassembleChatCompletionsEvents reassembles reasoning fields', async () => choices: [ { index: 0, - delta: { content: 'reply' }, + delta: { content: 'reply', reasoning_text: ' more', reasoning_opaque: 'enc' }, finish_reason: 'stop', }, ], @@ -182,11 +182,108 @@ test('reassembleChatCompletionsEvents reassembles reasoning fields', async () => const result = await reassembleChatCompletionsEvents(body); - assertEquals(result.choices[0].message.reasoning_text, 'think'); + assertEquals(result.choices[0].message.reasoning_text, 'think more'); assertEquals(result.choices[0].message.reasoning_opaque, 'enc'); assertEquals(result.choices[0].message.content, 'reply'); }); +test('reassembleChatCompletionsEvents maintains independent state for every choice index', async () => { + const body = makeEvents([ + { + data: { + id: 'cmpl_many', + object: 'chat.completion.chunk', + created: 4000, + model: 'gpt-many', + choices: [ + { + index: 1, + delta: { + role: 'assistant', + content: 'second ', + reasoning_opaque: 'second-old', + vendor_trace: 'b1', + tool_calls: [{ index: 0, id: 'call_b', type: 'function', function: { name: 'beta', arguments: '{"b"' } }], + }, + finish_reason: null, + content_filter_results: { hate: { filtered: false } }, + }, + { + index: 0, + delta: { + role: 'assistant', + content: 'first ', + reasoning_text: 'think ', + reasoning_opaque: 'first-old', + vendor_trace: 'a1', + reasoning_items: [{ type: 'reasoning', id: 'rs_a', summary: [] }], + tool_calls: [{ index: 1, id: 'call_a', type: 'function', function: { name: 'alpha', arguments: '{"a"' } }], + }, + finish_reason: null, + content_filter_results: { sexual: { filtered: false } }, + }, + ], + }, + }, + { + data: { + id: 'cmpl_many', + object: 'chat.completion.chunk', + created: 4000, + model: 'gpt-many', + choices: [ + { + index: 0, + delta: { + content: 'answer', + reasoning_text: 'again', + reasoning_opaque: 'first-final', + vendor_trace: 'a2', + tool_calls: [{ index: 1, function: { arguments: ':1}' } }], + }, + finish_reason: 'tool_calls', + }, + { + index: 1, + delta: { + content: 'answer', + reasoning_opaque: 'second-final', + vendor_trace: 'b2', + tool_calls: [{ index: 0, function: { arguments: ':2}' } }], + }, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 20, completion_tokens: 8, total_tokens: 28 }, + }, + }, + ]); + + const result = await reassembleChatCompletionsEvents(body) as ChatCompletionsResult & { + choices: Array; + }; + + assertEquals(result.choices.map(choice => choice.index), [0, 1]); + assertEquals(result.choices[0].message.content, 'first answer'); + assertEquals(result.choices[0].message.reasoning_text, 'think again'); + assertEquals(result.choices[0].message.reasoning_opaque, 'first-final'); + assertEquals(result.choices[0].message.reasoning_items, [{ type: 'reasoning', id: 'rs_a', summary: [] }]); + assertEquals(result.choices[0].message.tool_calls, [{ id: 'call_a', type: 'function', function: { name: 'alpha', arguments: '{"a":1}' } }]); + assertEquals(result.choices[0].message.vendor_trace, 'a1a2'); + assertEquals(result.choices[0].content_filter_results, { sexual: { filtered: false } }); + assertEquals(result.choices[0].finish_reason, 'tool_calls'); + assertEquals(result.choices[1].message.content, 'second answer'); + assertEquals(result.choices[1].message.reasoning_opaque, 'second-final'); + assertEquals(result.choices[1].message.tool_calls, [{ id: 'call_b', type: 'function', function: { name: 'beta', arguments: '{"b":2}' } }]); + assertEquals(result.choices[1].message.vendor_trace, 'b1b2'); + assertEquals(result.choices[1].content_filter_results, { hate: { filtered: false } }); + assertEquals(result.choices[1].finish_reason, 'stop'); + assertEquals(result.usage, { prompt_tokens: 20, completion_tokens: 8, total_tokens: 28 }); +}); + test('reassembleChatCompletionsEvents appends reasoning_items deltas in order', async () => { const body = makeEvents([ { diff --git a/packages/protocols/src/messages/reassemble.ts b/packages/protocols/src/messages/reassemble.ts index 05f0449d4..d4c6740c9 100644 --- a/packages/protocols/src/messages/reassemble.ts +++ b/packages/protocols/src/messages/reassemble.ts @@ -177,7 +177,7 @@ const applyBlockDelta = (block: MessagesBlockAccumulator | undefined, event: Ext return; case 'signature_delta': if (block.type !== 'thinking') return; - block.signature = `${block.signature ?? ''}${event.delta.signature ?? ''}`; + block.signature = event.delta.signature; return; } }; diff --git a/packages/protocols/src/messages/reassemble_test.ts b/packages/protocols/src/messages/reassemble_test.ts index a1360e11f..82998038c 100644 --- a/packages/protocols/src/messages/reassemble_test.ts +++ b/packages/protocols/src/messages/reassemble_test.ts @@ -340,6 +340,14 @@ test('reassembleMessagesEvents reassembles thinking blocks', async () => { delta: { type: 'thinking_delta', thinking: 'let me think' }, }, }, + { + event: 'content_block_delta', + data: { + type: 'content_block_delta', + index: 0, + delta: { type: 'signature_delta', signature: 'sig_old' }, + }, + }, { event: 'content_block_delta', data: { diff --git a/packages/translate/src/gemini-via-chat-completions/events.ts b/packages/translate/src/gemini-via-chat-completions/events.ts index 478b916d1..08cae52e1 100644 --- a/packages/translate/src/gemini-via-chat-completions/events.ts +++ b/packages/translate/src/gemini-via-chat-completions/events.ts @@ -1,4 +1,4 @@ -import { appendGeminiThoughtSignature, flushGeminiThoughtSignature, type GeminiThoughtSignatureState, parseStrictJsonObject, signGeminiPart } from '../shared/gemini-via/gemini.ts'; +import { flushGeminiThoughtSignature, type GeminiThoughtSignatureState, parseStrictJsonObject, setGeminiThoughtSignature, signGeminiPart } from '../shared/gemini-via/gemini.ts'; import { chatCompletionsErrorPayloadMessage } from '@floway-dev/protocols/chat-completions'; import type { ChatCompletionsStreamEvent, ChatCompletionsDelta } from '@floway-dev/protocols/chat-completions'; import { billableServiceTier, eventFrame, splitCacheWriteTokens, splitInclusiveInputTokens, splitInclusiveOutputTokens, USAGE_BILLING, type ProtocolFrame } from '@floway-dev/protocols/common'; @@ -137,7 +137,7 @@ const buildCandidate = (choice: ChatCompletionsStreamChoice, state: ChatCompleti } if (typeof delta.reasoning_opaque === 'string') { - appendGeminiThoughtSignature(state, delta.reasoning_opaque); + setGeminiThoughtSignature(state, delta.reasoning_opaque); } if (typeof delta.content === 'string') { diff --git a/packages/translate/src/gemini-via-chat-completions/events_test.ts b/packages/translate/src/gemini-via-chat-completions/events_test.ts index d30d09f11..372bce1ef 100644 --- a/packages/translate/src/gemini-via-chat-completions/events_test.ts +++ b/packages/translate/src/gemini-via-chat-completions/events_test.ts @@ -74,6 +74,7 @@ test('translateToSourceEvents maps text chunks and stop finish without emitting test('translateToSourceEvents maps reasoning text and attaches opaque signature to next action', async () => { const frames = await collect([ eventFrame(chunk({ role: 'assistant', reasoning_text: 'trace' })), + eventFrame(chunk({ reasoning_opaque: 'sig_old' })), eventFrame(chunk({ reasoning_opaque: 'sig_1' })), eventFrame(chunk({ content: 'answer' })), eventFrame(chunk({}, 'stop')), diff --git a/packages/translate/src/gemini-via-chat-completions/request.ts b/packages/translate/src/gemini-via-chat-completions/request.ts index 4a21a0d08..628c1b6e0 100644 --- a/packages/translate/src/gemini-via-chat-completions/request.ts +++ b/packages/translate/src/gemini-via-chat-completions/request.ts @@ -16,7 +16,7 @@ import { TranslatorInputError } from '../translator-input-error.ts'; import type { ChatCompletionsPayload, ChatCompletionsContentPart, ChatCompletionsMessage, ChatCompletionsTool, ChatCompletionsToolCall } from '@floway-dev/protocols/chat-completions'; import type { GeminiContent, GeminiPayload, GeminiGenerationConfig, GeminiPart } from '@floway-dev/protocols/gemini'; -const appendOpaque = (current: string | null, signature?: string): string | null => (typeof signature === 'string' ? `${current ?? ''}${signature}` : current); +const latestOpaque = (current: string | null, signature?: string): string | null => (typeof signature === 'string' ? signature : current); const inlineDataToContentPart = (part: GeminiPart): ChatCompletionsContentPart | null => { const url = geminiInlineDataUrl(part); @@ -56,7 +56,7 @@ const buildAssistantMessage = (content: GeminiContent, turnIndex: number, unmatc let reasoningOpaque: string | null = null; content.parts.forEach((part, partIndex) => { - reasoningOpaque = appendOpaque(reasoningOpaque, part.thoughtSignature); + reasoningOpaque = latestOpaque(reasoningOpaque, part.thoughtSignature); const kind = geminiPartKind(part); switch (kind) { diff --git a/packages/translate/src/gemini-via-chat-completions/request_test.ts b/packages/translate/src/gemini-via-chat-completions/request_test.ts index 5be6ab7af..1745dffe8 100644 --- a/packages/translate/src/gemini-via-chat-completions/request_test.ts +++ b/packages/translate/src/gemini-via-chat-completions/request_test.ts @@ -57,7 +57,7 @@ test('buildTargetRequest maps function calls, tool responses, and reasoning hist contents: [ { role: 'model', - parts: [{ text: 'private trace', thought: true, thoughtSignature: 'sig_1' }, { text: 'I will call a tool.' }, { functionCall: { name: 'lookup', args: { query: 'docs' } } }], + parts: [{ text: 'private trace', thought: true, thoughtSignature: 'sig_old' }, { text: 'I will call a tool.', thoughtSignature: 'sig_1' }, { functionCall: { name: 'lookup', args: { query: 'docs' } } }], }, { role: 'user', diff --git a/packages/translate/src/gemini-via-messages/events.ts b/packages/translate/src/gemini-via-messages/events.ts index 8e233099a..9ff8792f4 100644 --- a/packages/translate/src/gemini-via-messages/events.ts +++ b/packages/translate/src/gemini-via-messages/events.ts @@ -1,4 +1,4 @@ -import { appendGeminiThoughtSignature, flushGeminiThoughtSignature, type GeminiThoughtSignatureState, geminiCandidateEvent, parseStrictJsonObject, signGeminiPart } from '../shared/gemini-via/gemini.ts'; +import { flushGeminiThoughtSignature, type GeminiThoughtSignatureState, geminiCandidateEvent, parseStrictJsonObject, setGeminiThoughtSignature, signGeminiPart } from '../shared/gemini-via/gemini.ts'; import { billableServiceTier, eventFrame, splitInclusiveInputTokens, USAGE_BILLING, type ProtocolFrame } from '@floway-dev/protocols/common'; import type { GeminiFinishReason, GeminiStreamEvent, GeminiUsageMetadata } from '@floway-dev/protocols/gemini'; import { mergeMessagesUsageSnapshot, messagesUsageSnapshot, splitMessagesCacheCreationTokens, type MessagesStreamEvent, type MessagesUsageSnapshot } from '@floway-dev/protocols/messages'; @@ -108,7 +108,7 @@ export const translateToSourceEvents = async function* (frames: AsyncIterable 0) { diff --git a/packages/translate/src/gemini-via-messages/events_test.ts b/packages/translate/src/gemini-via-messages/events_test.ts index bd50de297..64af3f018 100644 --- a/packages/translate/src/gemini-via-messages/events_test.ts +++ b/packages/translate/src/gemini-via-messages/events_test.ts @@ -121,6 +121,11 @@ test('translateToSourceEvents maps thinking text and attaches signature to the n index: 0, delta: { type: 'thinking_delta', thinking: 'trace' }, }), + eventFrame({ + type: 'content_block_delta', + index: 0, + delta: { type: 'signature_delta', signature: 'sig_old' }, + }), eventFrame({ type: 'content_block_delta', index: 0, diff --git a/packages/translate/src/messages-via-chat-completions/events.ts b/packages/translate/src/messages-via-chat-completions/events.ts index 775b0fede..2e54d5926 100644 --- a/packages/translate/src/messages-via-chat-completions/events.ts +++ b/packages/translate/src/messages-via-chat-completions/events.ts @@ -143,7 +143,7 @@ const attachOpaqueToOpenThinkingBlock = (state: ChatCompletionsToMessagesStreamS return false; } - state.pendingThinkingSignature = (state.pendingThinkingSignature ?? '') + state.pendingReasoningOpaque; + state.pendingThinkingSignature = state.pendingReasoningOpaque; state.pendingReasoningOpaque = undefined; return true; }; @@ -232,12 +232,12 @@ const handleReasoningDelta = (delta: ChatCompletionsStreamDelta, state: ChatComp } if (state.openBlock === 'thinking') { - state.pendingThinkingSignature = (state.pendingThinkingSignature ?? '') + delta.reasoning_opaque; + state.pendingThinkingSignature = delta.reasoning_opaque; emitPendingReasoningAndDeferred(state, events); return; } - state.pendingReasoningOpaque = (state.pendingReasoningOpaque ?? '') + delta.reasoning_opaque; + state.pendingReasoningOpaque = delta.reasoning_opaque; }; const emitToolCallsDelta = (toolCalls: ChatCompletionsStreamToolCalls, state: ChatCompletionsToMessagesStreamState, events: MessagesStreamEvent[]): void => { diff --git a/packages/translate/src/messages-via-chat-completions/events_test.ts b/packages/translate/src/messages-via-chat-completions/events_test.ts index c90d946d9..adb39dd17 100644 --- a/packages/translate/src/messages-via-chat-completions/events_test.ts +++ b/packages/translate/src/messages-via-chat-completions/events_test.ts @@ -28,7 +28,8 @@ const usageChunk = (): ChatCompletionsStreamEvent => ({ test('translateChatCompletionsChunkToMessagesEvents emits opaque-only reasoning as redacted_thinking at finish', () => { const state = createChatCompletionsToMessagesStreamState(); const events = [ - ...translateChatCompletionsChunkToMessagesEvents(chunk({ role: 'assistant', reasoning_opaque: 'enc_only' }), state), + ...translateChatCompletionsChunkToMessagesEvents(chunk({ role: 'assistant', reasoning_opaque: 'enc_old' }), state), + ...translateChatCompletionsChunkToMessagesEvents(chunk({ reasoning_opaque: 'enc_only' }), state), ...translateChatCompletionsChunkToMessagesEvents(chunk({}, 'stop'), state), ]; diff --git a/packages/translate/src/shared/gemini-via/gemini.ts b/packages/translate/src/shared/gemini-via/gemini.ts index 0ff026a62..3cb8819e5 100644 --- a/packages/translate/src/shared/gemini-via/gemini.ts +++ b/packages/translate/src/shared/gemini-via/gemini.ts @@ -185,8 +185,8 @@ export interface GeminiThoughtSignatureState { pendingThoughtSignature?: string; } -export const appendGeminiThoughtSignature = (state: GeminiThoughtSignatureState, signature: string): void => { - state.pendingThoughtSignature = `${state.pendingThoughtSignature ?? ''}${signature}`; +export const setGeminiThoughtSignature = (state: GeminiThoughtSignatureState, signature: string): void => { + state.pendingThoughtSignature = signature; }; export const signGeminiPart = (state: GeminiThoughtSignatureState, part: GeminiPart): GeminiPart => { From b648cb472fcfcb5ede16e24f74981d6d578a7e25 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:52:35 +0800 Subject: [PATCH 006/157] feat: add protocol affinity egress transforms Wrap or synthesize client-carried affinity at each source protocol boundary while keeping visible stream content ahead of delayed opaque-carrier encoding. Preserve per-choice and per-candidate semantics and reuse Responses replacements across lifecycle snapshots. --- .../chat/affinity/affinity-egress.ts | 10 ++ .../chat/affinity/chat-completions-egress.ts | 80 ++++++++++ .../affinity/chat-completions-egress_test.ts | 101 +++++++++++++ .../data-plane/chat/affinity/gemini-egress.ts | 100 ++++++++++++ .../chat/affinity/gemini-egress_test.ts | 91 +++++++++++ .../chat/affinity/messages-egress.ts | 106 +++++++++++++ .../chat/affinity/messages-egress_test.ts | 104 +++++++++++++ .../chat/affinity/responses-egress.ts | 120 +++++++++++++++ .../chat/affinity/responses-egress_test.ts | 143 ++++++++++++++++++ 9 files changed, 855 insertions(+) create mode 100644 packages/gateway/src/data-plane/chat/affinity/affinity-egress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/chat-completions-egress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/chat-completions-egress_test.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/gemini-egress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/gemini-egress_test.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/messages-egress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/messages-egress_test.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/responses-egress.ts create mode 100644 packages/gateway/src/data-plane/chat/affinity/responses-egress_test.ts diff --git a/packages/gateway/src/data-plane/chat/affinity/affinity-egress.ts b/packages/gateway/src/data-plane/chat/affinity/affinity-egress.ts new file mode 100644 index 000000000..d326a6f54 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/affinity-egress.ts @@ -0,0 +1,10 @@ +import type { AffinityTarget } from './types.ts'; + +export interface AffinityEgressCodec { + wrap(value: string | undefined, affinity: AffinityTarget): Promise; +} + +export interface AffinityEgressOptions { + readonly codec: AffinityEgressCodec; + readonly affinity: AffinityTarget; +} diff --git a/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress.ts b/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress.ts new file mode 100644 index 000000000..bf631c8c2 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress.ts @@ -0,0 +1,80 @@ +import type { AffinityEgressOptions } from './affinity-egress.ts'; +import type { ChatCompletionsDelta, ChatCompletionsStreamEvent } from '@floway-dev/protocols/chat-completions'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; + +interface ChoiceState { + opaque?: string; + finished: boolean; +} + +type StreamingChoice = ChatCompletionsStreamEvent['choices'][number]; + +const withoutOpaque = (delta: ChatCompletionsDelta): ChatCompletionsDelta => { + const { reasoning_opaque: _opaque, ...visible } = delta; + return visible; +}; + +const eventWithChoices = ( + event: ChatCompletionsStreamEvent, + choices: StreamingChoice[], + includeUsage: boolean, +): ChatCompletionsStreamEvent => { + if (includeUsage || event.usage === undefined) return { ...event, choices }; + const { usage: _usage, ...rest } = event; + return { ...rest, choices }; +}; + +export const wrapChatCompletionsAffinityEgress = async function* ( + frames: AsyncIterable>, + options: AffinityEgressOptions, +): AsyncGenerator> { + const choices = new Map(); + + for await (const frame of frames) { + if (frame.type !== 'event') { + yield frame; + continue; + } + + const visibleChoices: StreamingChoice[] = []; + const finishingChoices: Array<{ choice: StreamingChoice; state: ChoiceState }> = []; + + for (const choice of frame.event.choices) { + const state = choices.get(choice.index) ?? { finished: false }; + if (state.finished) throw new Error(`Chat Completions choice ${choice.index} emitted data after its finish_reason`); + choices.set(choice.index, state); + + if (typeof choice.delta.reasoning_opaque === 'string') state.opaque = choice.delta.reasoning_opaque; + const delta = withoutOpaque(choice.delta); + + if (choice.finish_reason === null) { + if (Object.keys(delta).length > 0) visibleChoices.push({ ...choice, delta }); + continue; + } + + if (Object.keys(delta).length > 0) { + visibleChoices.push({ ...choice, delta, finish_reason: null }); + } + finishingChoices.push({ choice, state }); + } + + if (visibleChoices.length > 0 || frame.event.choices.length === 0) { + yield eventFrame(eventWithChoices(frame.event, visibleChoices, finishingChoices.length === 0)); + } + + if (finishingChoices.length === 0) continue; + + const wrappedChoices = await Promise.all(finishingChoices.map(async ({ choice, state }) => ({ + ...choice, + delta: { reasoning_opaque: await options.codec.wrap(state.opaque, options.affinity) }, + finish_reason: null, + }))); + yield eventFrame(eventWithChoices(frame.event, wrappedChoices, false)); + + const finishedChoices = finishingChoices.map(({ choice, state }) => { + state.finished = true; + return { ...choice, delta: {}, finish_reason: choice.finish_reason }; + }); + yield eventFrame(eventWithChoices(frame.event, finishedChoices, true)); + } +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress_test.ts b/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress_test.ts new file mode 100644 index 000000000..afa600849 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/chat-completions-egress_test.ts @@ -0,0 +1,101 @@ +import { describe, expect, test } from 'vitest'; + +import type { AffinityEgressCodec } from './affinity-egress.ts'; +import { wrapChatCompletionsAffinityEgress } from './chat-completions-egress.ts'; +import type { AffinityTarget } from './types.ts'; +import type { ChatCompletionsStreamEvent } from '@floway-dev/protocols/chat-completions'; +import { doneFrame, eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; + +const affinity: AffinityTarget = { + mode: 'prefer', + upstreamId: 'up-a', + upstreamRevision: 'rev-a', + modelId: 'model-a', + rulesPresent: false, +}; + +const chunk = ( + choices: ChatCompletionsStreamEvent['choices'], +): ChatCompletionsStreamEvent => ({ + id: 'chatcmpl_1', + object: 'chat.completion.chunk', + created: 1, + model: 'model-a', + choices, +}); + +const frames = async function* (values: ProtocolFrame[]) { + yield* values; +}; + +class DelayedCodec implements AffinityEgressCodec { + readonly calls: Array<{ value: string | undefined; resolve: (value: string) => void }> = []; + + wrap(value: string | undefined): Promise { + return new Promise(resolve => this.calls.push({ value, resolve })); + } +} + +const immediateCodec: AffinityEgressCodec = { + wrap: async value => `wrapped:${value ?? 'synthetic'}`, +}; + +describe('Chat Completions affinity egress', () => { + test('forwards visible final data before wrapping the last opaque snapshot', async () => { + const codec = new DelayedCodec(); + const output = wrapChatCompletionsAffinityEgress(frames([ + eventFrame(chunk([{ index: 0, delta: { reasoning_opaque: 'first' }, finish_reason: null }])), + eventFrame(chunk([{ + index: 0, + delta: { content: 'visible', reasoning_text: 'thinking', reasoning_opaque: 'latest' }, + finish_reason: 'stop', + }])), + doneFrame(), + ]), { codec, affinity })[Symbol.asyncIterator](); + + const visible = await output.next(); + expect(visible.value).toEqual(eventFrame(chunk([{ + index: 0, + delta: { content: 'visible', reasoning_text: 'thinking' }, + finish_reason: null, + }]))); + expect(codec.calls).toHaveLength(0); + + const wrappedPending = output.next(); + await Promise.resolve(); + expect(codec.calls.map(call => call.value)).toEqual(['latest']); + codec.calls[0].resolve('wrapped-latest'); + expect((await wrappedPending).value).toEqual(eventFrame(chunk([{ + index: 0, + delta: { reasoning_opaque: 'wrapped-latest' }, + finish_reason: null, + }]))); + + expect((await output.next()).value).toEqual(eventFrame(chunk([{ + index: 0, + delta: {}, + finish_reason: 'stop', + }]))); + expect((await output.next()).value).toEqual(doneFrame()); + }); + + test('wraps or synthesizes a carrier independently for every finishing choice', async () => { + const output: ProtocolFrame[] = []; + for await (const frame of wrapChatCompletionsAffinityEgress(frames([ + eventFrame(chunk([ + { index: 0, delta: { reasoning_opaque: 'opaque' }, finish_reason: 'stop' }, + { index: 1, delta: {}, finish_reason: 'length' }, + ])), + doneFrame(), + ]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output[0]).toEqual(eventFrame(chunk([ + { index: 0, delta: { reasoning_opaque: 'wrapped:opaque' }, finish_reason: null }, + { index: 1, delta: { reasoning_opaque: 'wrapped:synthetic' }, finish_reason: null }, + ]))); + expect(output[1]).toEqual(eventFrame(chunk([ + { index: 0, delta: {}, finish_reason: 'stop' }, + { index: 1, delta: {}, finish_reason: 'length' }, + ]))); + }); +}); diff --git a/packages/gateway/src/data-plane/chat/affinity/gemini-egress.ts b/packages/gateway/src/data-plane/chat/affinity/gemini-egress.ts new file mode 100644 index 000000000..1e0e1b1e3 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/gemini-egress.ts @@ -0,0 +1,100 @@ +import type { AffinityEgressOptions } from './affinity-egress.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { GeminiCandidate, GeminiPart, GeminiStreamEvent } from '@floway-dev/protocols/gemini'; + +interface CandidateState { + sawCarrier: boolean; + finished: boolean; +} + +interface DeferredCandidate { + readonly candidate: GeminiCandidate; + readonly carrierValues: string[]; + readonly state: CandidateState; +} + +const withoutThoughtSignature = (part: GeminiPart): GeminiPart => { + const { thoughtSignature: _signature, ...visible } = part; + return visible; +}; + +const hasPartData = (part: GeminiPart): boolean => Object.keys(part).length > 0; + +const eventWithoutUsage = (event: GeminiStreamEvent): GeminiStreamEvent => { + if ('error' in event || event.usageMetadata === undefined) return event; + const { usageMetadata: _usage, ...rest } = event; + return rest; +}; + +export const wrapGeminiAffinityEgress = async function* ( + frames: AsyncIterable>, + options: AffinityEgressOptions, +): AsyncGenerator> { + const states = new Map(); + + for await (const frame of frames) { + if (frame.type !== 'event') { + yield frame; + continue; + } + if ('error' in frame.event) { + yield frame; + return; + } + + const visibleCandidates: GeminiCandidate[] = []; + const deferredCandidates: DeferredCandidate[] = []; + + for (const candidate of frame.event.candidates ?? []) { + const state = states.get(candidate.index) ?? { sawCarrier: false, finished: false }; + if (state.finished) throw new Error(`Gemini candidate ${candidate.index} emitted data after its finishReason`); + states.set(candidate.index, state); + + const carrierValues: string[] = []; + const visibleParts: GeminiPart[] = []; + for (const part of candidate.content.parts) { + if (typeof part.thoughtSignature === 'string') carrierValues.push(part.thoughtSignature); + const visible = withoutThoughtSignature(part); + if (hasPartData(visible)) visibleParts.push(visible); + } + + const needsSynthetic = candidate.finishReason !== undefined && !state.sawCarrier && carrierValues.length === 0; + if (carrierValues.length > 0 || needsSynthetic) { + deferredCandidates.push({ candidate, carrierValues, state }); + if (visibleParts.length > 0) { + const { finishReason: _finishReason, ...rest } = candidate; + visibleCandidates.push({ ...rest, content: { ...candidate.content, parts: visibleParts } }); + } + continue; + } + + visibleCandidates.push(candidate); + if (candidate.finishReason !== undefined) state.finished = true; + } + + if (visibleCandidates.length > 0 || (frame.event.candidates?.length ?? 0) === 0) { + const visibleEvent = deferredCandidates.length > 0 ? eventWithoutUsage(frame.event) : frame.event; + yield eventFrame({ ...visibleEvent, candidates: visibleCandidates }); + } + + if (deferredCandidates.length === 0) continue; + + const wrappedCandidates = await Promise.all(deferredCandidates.map(async ({ candidate, carrierValues, state }) => { + const values = carrierValues.length > 0 ? carrierValues : [undefined]; + const parts = await Promise.all(values.map(async value => ({ + thoughtSignature: await options.codec.wrap(value, options.affinity), + }))); + state.sawCarrier = true; + if (candidate.finishReason !== undefined) state.finished = true; + return { + index: candidate.index, + content: { + ...(candidate.content.role !== undefined ? { role: candidate.content.role } : {}), + parts, + }, + ...(candidate.finishReason !== undefined ? { finishReason: candidate.finishReason } : {}), + } satisfies GeminiCandidate; + })); + yield eventFrame({ ...frame.event, candidates: wrappedCandidates }); + } +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/gemini-egress_test.ts b/packages/gateway/src/data-plane/chat/affinity/gemini-egress_test.ts new file mode 100644 index 000000000..4672708b9 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/gemini-egress_test.ts @@ -0,0 +1,91 @@ +import { describe, expect, test } from 'vitest'; + +import type { AffinityEgressCodec } from './affinity-egress.ts'; +import { wrapGeminiAffinityEgress } from './gemini-egress.ts'; +import type { AffinityTarget } from './types.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { GeminiStreamEvent } from '@floway-dev/protocols/gemini'; + +const affinity: AffinityTarget = { + mode: 'prefer', + upstreamId: 'up-a', + upstreamRevision: 'rev-a', + modelId: 'model-a', + rulesPresent: false, +}; + +const frames = async function* (values: ProtocolFrame[]) { + yield* values; +}; + +class DelayedCodec implements AffinityEgressCodec { + readonly calls: Array<{ value: string | undefined; resolve: (value: string) => void }> = []; + + wrap(value: string | undefined): Promise { + return new Promise(resolve => this.calls.push({ value, resolve })); + } +} + +const immediateCodec: AffinityEgressCodec = { + wrap: async value => `wrapped:${value ?? 'synthetic'}`, +}; + +describe('Gemini affinity egress', () => { + test('emits visible part data before wrapping a thoughtSignature from the same part', async () => { + const codec = new DelayedCodec(); + const input: GeminiStreamEvent = { + candidates: [{ + index: 0, + content: { role: 'model', parts: [{ text: 'visible', thoughtSignature: 'opaque' }] }, + finishReason: 'STOP', + }], + usageMetadata: { totalTokenCount: 2 }, + }; + const output = wrapGeminiAffinityEgress(frames([eventFrame(input)]), { codec, affinity })[Symbol.asyncIterator](); + + expect((await output.next()).value).toEqual(eventFrame({ + candidates: [{ index: 0, content: { role: 'model', parts: [{ text: 'visible' }] } }], + })); + expect(codec.calls).toHaveLength(0); + + const carrierPending = output.next(); + await Promise.resolve(); + expect(codec.calls.map(call => call.value)).toEqual(['opaque']); + codec.calls[0].resolve('wrapped-opaque'); + expect((await carrierPending).value).toEqual(eventFrame({ + candidates: [{ + index: 0, + content: { role: 'model', parts: [{ thoughtSignature: 'wrapped-opaque' }] }, + finishReason: 'STOP', + }], + usageMetadata: { totalTokenCount: 2 }, + })); + }); + + test('synthesizes a signature-only part for each finishing candidate without one', async () => { + const output: ProtocolFrame[] = []; + for await (const frame of wrapGeminiAffinityEgress(frames([eventFrame({ + candidates: [ + { index: 0, content: { role: 'model', parts: [{ text: 'a' }] }, finishReason: 'STOP' }, + { index: 1, content: { role: 'model', parts: [{ functionCall: { name: 'tool', args: {} } }] }, finishReason: 'MAX_TOKENS' }, + ], + })]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output[0]).toMatchObject({ + event: { + candidates: [ + { index: 0, content: { parts: [{ text: 'a' }] } }, + { index: 1, content: { parts: [{ functionCall: { name: 'tool' } }] } }, + ], + }, + }); + expect(output[1]).toMatchObject({ + event: { + candidates: [ + { index: 0, content: { parts: [{ thoughtSignature: 'wrapped:synthetic' }] }, finishReason: 'STOP' }, + { index: 1, content: { parts: [{ thoughtSignature: 'wrapped:synthetic' }] }, finishReason: 'MAX_TOKENS' }, + ], + }, + }); + }); +}); diff --git a/packages/gateway/src/data-plane/chat/affinity/messages-egress.ts b/packages/gateway/src/data-plane/chat/affinity/messages-egress.ts new file mode 100644 index 000000000..f71156d8e --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/messages-egress.ts @@ -0,0 +1,106 @@ +import type { AffinityEgressOptions } from './affinity-egress.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { MessagesStreamEvent } from '@floway-dev/protocols/messages'; + +interface OpenBlock { + readonly type: string; + signature?: string; +} + +export const wrapMessagesAffinityEgress = async function* ( + frames: AsyncIterable>, + options: AffinityEgressOptions, +): AsyncGenerator> { + const openBlocks = new Map(); + let nextBlockIndex = 0; + let sawCarrier = false; + let syntheticEmitted = false; + + const syntheticEvents = async (): Promise => { + if (sawCarrier || syntheticEmitted) return []; + syntheticEmitted = true; + const index = nextBlockIndex++; + return [ + { + type: 'content_block_start', + index, + content_block: { + type: 'redacted_thinking', + data: await options.codec.wrap(undefined, options.affinity), + }, + }, + { type: 'content_block_stop', index }, + ]; + }; + + for await (const frame of frames) { + if (frame.type !== 'event') { + yield frame; + continue; + } + + const event = frame.event; + if (event.type === 'content_block_start') { + nextBlockIndex = Math.max(nextBlockIndex, event.index + 1); + if (openBlocks.has(event.index)) throw new Error(`Messages content block ${event.index} started twice`); + openBlocks.set(event.index, { type: event.content_block.type }); + if (event.content_block.type !== 'redacted_thinking') { + yield frame; + continue; + } + + sawCarrier = true; + yield eventFrame({ + ...event, + content_block: { + ...event.content_block, + data: await options.codec.wrap(event.content_block.data, options.affinity), + }, + }); + continue; + } + + if (event.type === 'content_block_delta' && event.delta.type === 'signature_delta') { + const block = openBlocks.get(event.index); + if (block?.type !== 'thinking') throw new Error(`Messages signature_delta targeted non-thinking block ${event.index}`); + block.signature = event.delta.signature; + continue; + } + + if (event.type === 'content_block_stop') { + const block = openBlocks.get(event.index); + if (block === undefined) throw new Error(`Messages content block ${event.index} stopped before it started`); + if (block.signature !== undefined) { + sawCarrier = true; + yield eventFrame({ + type: 'content_block_delta', + index: event.index, + delta: { + type: 'signature_delta', + signature: await options.codec.wrap(block.signature, options.affinity), + }, + }); + } + openBlocks.delete(event.index); + yield frame; + continue; + } + + if (event.type === 'message_delta' && event.delta.stop_reason != null) { + if (openBlocks.size > 0) throw new Error('Messages terminal message_delta arrived with an open content block'); + for (const synthetic of await syntheticEvents()) yield eventFrame(synthetic); + yield frame; + continue; + } + + if (event.type === 'message_stop') { + if (openBlocks.size > 0) throw new Error('Messages message_stop arrived with an open content block'); + for (const synthetic of await syntheticEvents()) yield eventFrame(synthetic); + yield frame; + return; + } + + yield frame; + if (event.type === 'error') return; + } +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/messages-egress_test.ts b/packages/gateway/src/data-plane/chat/affinity/messages-egress_test.ts new file mode 100644 index 000000000..f208ef2fb --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/messages-egress_test.ts @@ -0,0 +1,104 @@ +import { describe, expect, test } from 'vitest'; + +import type { AffinityEgressCodec } from './affinity-egress.ts'; +import { wrapMessagesAffinityEgress } from './messages-egress.ts'; +import type { AffinityTarget } from './types.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { MessagesStreamEvent } from '@floway-dev/protocols/messages'; + +const affinity: AffinityTarget = { + mode: 'prefer', + upstreamId: 'up-a', + upstreamRevision: 'rev-a', + modelId: 'model-a', + rulesPresent: false, +}; + +const frames = async function* (values: ProtocolFrame[]) { + yield* values; +}; + +class DelayedCodec implements AffinityEgressCodec { + readonly calls: Array<{ value: string | undefined; resolve: (value: string) => void }> = []; + + wrap(value: string | undefined): Promise { + return new Promise(resolve => this.calls.push({ value, resolve })); + } +} + +const immediateCodec: AffinityEgressCodec = { + wrap: async value => `wrapped:${value ?? 'synthetic'}`, +}; + +describe('Messages affinity egress', () => { + test('streams readable thinking and wraps the latest signature snapshot at block stop', async () => { + const codec = new DelayedCodec(); + const output = wrapMessagesAffinityEgress(frames([ + eventFrame({ type: 'content_block_start', index: 0, content_block: { type: 'thinking', thinking: '' } }), + eventFrame({ type: 'content_block_delta', index: 0, delta: { type: 'thinking_delta', thinking: 'visible' } }), + eventFrame({ type: 'content_block_delta', index: 0, delta: { type: 'signature_delta', signature: 'first' } }), + eventFrame({ type: 'content_block_delta', index: 0, delta: { type: 'signature_delta', signature: 'latest' } }), + eventFrame({ type: 'content_block_stop', index: 0 }), + eventFrame({ type: 'message_delta', delta: { stop_reason: 'end_turn' } }), + eventFrame({ type: 'message_stop' }), + ]), { codec, affinity })[Symbol.asyncIterator](); + + expect((await output.next()).value).toMatchObject({ event: { type: 'content_block_start' } }); + expect((await output.next()).value).toEqual(eventFrame({ + type: 'content_block_delta', + index: 0, + delta: { type: 'thinking_delta', thinking: 'visible' }, + })); + expect(codec.calls).toHaveLength(0); + + const signaturePending = output.next(); + await Promise.resolve(); + expect(codec.calls.map(call => call.value)).toEqual(['latest']); + codec.calls[0].resolve('wrapped-latest'); + expect((await signaturePending).value).toEqual(eventFrame({ + type: 'content_block_delta', + index: 0, + delta: { type: 'signature_delta', signature: 'wrapped-latest' }, + })); + expect((await output.next()).value).toEqual(eventFrame({ type: 'content_block_stop', index: 0 })); + expect((await output.next()).value).toMatchObject({ event: { type: 'message_delta' } }); + }); + + test('wraps redacted data inline and does not inject another carrier', async () => { + const output: ProtocolFrame[] = []; + for await (const frame of wrapMessagesAffinityEgress(frames([ + eventFrame({ type: 'content_block_start', index: 2, content_block: { type: 'redacted_thinking', data: 'opaque' } }), + eventFrame({ type: 'content_block_stop', index: 2 }), + eventFrame({ type: 'message_delta', delta: { stop_reason: 'end_turn' } }), + eventFrame({ type: 'message_stop' }), + ]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output[0]).toEqual(eventFrame({ + type: 'content_block_start', + index: 2, + content_block: { type: 'redacted_thinking', data: 'wrapped:opaque' }, + })); + expect(output.filter(frame => frame.type === 'event' && frame.event.type === 'content_block_start')).toHaveLength(1); + }); + + test('injects a synthetic redacted block before the terminal message_delta', async () => { + const output: ProtocolFrame[] = []; + for await (const frame of wrapMessagesAffinityEgress(frames([ + eventFrame({ type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }), + eventFrame({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'answer' } }), + eventFrame({ type: 'content_block_stop', index: 0 }), + eventFrame({ type: 'message_delta', delta: { stop_reason: 'end_turn' } }), + eventFrame({ type: 'message_stop' }), + ]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output.slice(3, 6)).toEqual([ + eventFrame({ + type: 'content_block_start', + index: 1, + content_block: { type: 'redacted_thinking', data: 'wrapped:synthetic' }, + }), + eventFrame({ type: 'content_block_stop', index: 1 }), + eventFrame({ type: 'message_delta', delta: { stop_reason: 'end_turn' } }), + ]); + }); +}); diff --git a/packages/gateway/src/data-plane/chat/affinity/responses-egress.ts b/packages/gateway/src/data-plane/chat/affinity/responses-egress.ts new file mode 100644 index 000000000..83263d468 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/responses-egress.ts @@ -0,0 +1,120 @@ +import type { AffinityEgressOptions } from './affinity-egress.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { ResponsesOutputItem, ResponsesOutputReasoning, ResponsesResult, ResponsesStreamEvent } from '@floway-dev/protocols/responses'; + +const randomReasoningId = (): string => { + const bytes = crypto.getRandomValues(new Uint8Array(16)); + let hex = ''; + for (const byte of bytes) hex += byte.toString(16).padStart(2, '0'); + return `rs_affinity_${hex}`; +}; + +const encryptedContentSlots = (item: ResponsesOutputItem): Array<{ key: string; value: string }> => { + const slots: Array<{ key: string; value: string }> = []; + const record = item as unknown as Record; + if (typeof record.encrypted_content === 'string') { + slots.push({ key: 'encrypted_content', value: record.encrypted_content }); + } + if (Array.isArray(record.content)) { + record.content.forEach((part, index) => { + if (!part || typeof part !== 'object') return; + const content = part as Record; + if (content.type === 'encrypted_content' && typeof content.encrypted_content === 'string') { + slots.push({ key: `content.${index}.encrypted_content`, value: content.encrypted_content }); + } + }); + } + return slots; +}; + +const replaceEncryptedContent = ( + item: ResponsesOutputItem, + replacements: ReadonlyMap, +): ResponsesOutputItem => { + const record = item as unknown as Record; + const topLevel = replacements.get('encrypted_content'); + const content = Array.isArray(record.content) + ? record.content.map((part, index) => { + const replacement = replacements.get(`content.${index}.encrypted_content`); + return replacement === undefined || !part || typeof part !== 'object' + ? part + : { ...(part as Record), encrypted_content: replacement }; + }) + : undefined; + return { + ...item, + ...(topLevel !== undefined ? { encrypted_content: topLevel } : {}), + ...(content !== undefined ? { content } : {}), + } as ResponsesOutputItem; +}; + +export const wrapResponsesAffinityEgress = async function* ( + frames: AsyncIterable>, + options: AffinityEgressOptions, +): AsyncGenerator> { + const wrapped = new Map>(); + let sawCarrier = false; + + const wrapItem = async (item: ResponsesOutputItem, outputIndex: number): Promise => { + const slots = encryptedContentSlots(item); + if (slots.length === 0) return item; + sawCarrier = true; + const replacements = new Map(); + await Promise.all(slots.map(async slot => { + const cacheKey = `${outputIndex}\0${slot.key}\0${slot.value}`; + let replacement = wrapped.get(cacheKey); + if (replacement === undefined) { + replacement = options.codec.wrap(slot.value, options.affinity); + wrapped.set(cacheKey, replacement); + } + replacements.set(slot.key, await replacement); + })); + return replaceEncryptedContent(item, replacements); + }; + + const wrapResult = async (response: ResponsesResult): Promise => ({ + ...response, + output: await Promise.all(response.output.map(async (item, index) => await wrapItem(item, index))), + }); + + for await (const frame of frames) { + if (frame.type !== 'event') { + yield frame; + continue; + } + + const event = frame.event; + if (event.type === 'response.output_item.added' || event.type === 'response.output_item.done') { + yield eventFrame({ ...event, item: await wrapItem(event.item, event.output_index) }); + continue; + } + + if (event.type === 'response.completed' || event.type === 'response.incomplete') { + const response = await wrapResult(event.response); + if (sawCarrier) { + yield eventFrame({ ...event, response }); + return; + } + + const outputIndex = response.output.length; + const item: ResponsesOutputReasoning = { + type: 'reasoning', + id: randomReasoningId(), + summary: [], + encrypted_content: await options.codec.wrap(undefined, options.affinity), + }; + yield eventFrame({ type: 'response.output_item.added', output_index: outputIndex, item }); + yield eventFrame({ type: 'response.output_item.done', output_index: outputIndex, item }); + yield eventFrame({ ...event, response: { ...response, output: [...response.output, item] } }); + return; + } + + if (event.type === 'response.failed') { + yield eventFrame({ ...event, response: await wrapResult(event.response) }); + return; + } + + yield frame; + if (event.type === 'error') return; + } +}; diff --git a/packages/gateway/src/data-plane/chat/affinity/responses-egress_test.ts b/packages/gateway/src/data-plane/chat/affinity/responses-egress_test.ts new file mode 100644 index 000000000..4fd902a45 --- /dev/null +++ b/packages/gateway/src/data-plane/chat/affinity/responses-egress_test.ts @@ -0,0 +1,143 @@ +import { describe, expect, test } from 'vitest'; + +import type { AffinityEgressCodec } from './affinity-egress.ts'; +import { wrapResponsesAffinityEgress } from './responses-egress.ts'; +import type { AffinityTarget } from './types.ts'; +import { eventFrame, type ProtocolFrame } from '@floway-dev/protocols/common'; +import type { ResponsesOutputReasoning, ResponsesResult, ResponsesStreamEvent } from '@floway-dev/protocols/responses'; + +const affinity: AffinityTarget = { + mode: 'prefer', + upstreamId: 'up-a', + upstreamRevision: 'rev-a', + modelId: 'model-a', + rulesPresent: false, +}; + +const frames = async function* (values: ProtocolFrame[]) { + yield* values; +}; + +class DelayedCodec implements AffinityEgressCodec { + readonly calls: Array<{ value: string | undefined; resolve: (value: string) => void }> = []; + + wrap(value: string | undefined): Promise { + return new Promise(resolve => this.calls.push({ value, resolve })); + } +} + +const immediateCodec: AffinityEgressCodec = { + wrap: async value => `wrapped:${value ?? 'synthetic'}`, +}; + +const response = (output: ResponsesResult['output'], status: ResponsesResult['status'] = 'completed'): ResponsesResult => ({ + id: 'resp_1', + object: 'response', + model: 'model-a', + output, + status, + error: null, + incomplete_details: null, +}); + +describe('Responses affinity egress', () => { + test('passes reasoning summary deltas before wrapping and reuses one exact replacement', async () => { + const codec = new DelayedCodec(); + const item: ResponsesOutputReasoning = { + type: 'reasoning', + id: 'rs_1', + summary: [{ type: 'summary_text', text: 'visible' }], + encrypted_content: 'opaque', + }; + const output = wrapResponsesAffinityEgress(frames([ + eventFrame({ + type: 'response.reasoning_summary_text.delta', + item_id: 'rs_1', + output_index: 0, + summary_index: 0, + delta: 'visible', + }), + eventFrame({ type: 'response.output_item.done', output_index: 0, item }), + eventFrame({ type: 'response.completed', response: response([item]) }), + ]), { codec, affinity })[Symbol.asyncIterator](); + + expect((await output.next()).value).toMatchObject({ + event: { type: 'response.reasoning_summary_text.delta', delta: 'visible' }, + }); + expect(codec.calls).toHaveLength(0); + + const donePending = output.next(); + await Promise.resolve(); + expect(codec.calls.map(call => call.value)).toEqual(['opaque']); + codec.calls[0].resolve('wrapped-opaque'); + expect((await donePending).value).toMatchObject({ + event: { item: { encrypted_content: 'wrapped-opaque' } }, + }); + + expect((await output.next()).value).toMatchObject({ + event: { response: { output: [{ encrypted_content: 'wrapped-opaque' }] } }, + }); + expect(codec.calls).toHaveLength(1); + }); + + test('caches replacements across added, done, and terminal snapshots', async () => { + const calls: Array = []; + const codec: AffinityEgressCodec = { + wrap: async value => { + calls.push(value); + return `wrapped:${value}`; + }, + }; + const item: ResponsesOutputReasoning = { type: 'reasoning', id: 'rs_1', summary: [], encrypted_content: 'opaque' }; + const output: ProtocolFrame[] = []; + for await (const frame of wrapResponsesAffinityEgress(frames([ + eventFrame({ type: 'response.output_item.added', output_index: 0, item }), + eventFrame({ type: 'response.output_item.done', output_index: 0, item }), + eventFrame({ type: 'response.completed', response: response([item]) }), + ]), { codec, affinity })) output.push(frame); + + expect(calls).toEqual(['opaque']); + expect(output).toHaveLength(3); + for (const frame of output) expect(JSON.stringify(frame)).toContain('wrapped:opaque'); + }); + + test('injects one synthetic reasoning lifecycle before a carrier-free terminal response', async () => { + const message = { + type: 'message' as const, + id: 'msg_1', + role: 'assistant' as const, + status: 'completed', + content: [{ type: 'output_text' as const, text: 'answer' }], + }; + const output: ProtocolFrame[] = []; + for await (const frame of wrapResponsesAffinityEgress(frames([ + eventFrame({ type: 'response.completed', response: response([message]) }), + ]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output.map(frame => frame.type === 'event' ? frame.event.type : frame.type)).toEqual([ + 'response.output_item.added', + 'response.output_item.done', + 'response.completed', + ]); + const added = output[0]; + const done = output[1]; + if (added.type !== 'event' || added.event.type !== 'response.output_item.added') throw new Error('Expected added event'); + if (done.type !== 'event' || done.event.type !== 'response.output_item.done') throw new Error('Expected done event'); + expect(added.event.item).toEqual(done.event.item); + expect(added.event.item).toMatchObject({ + type: 'reasoning', + summary: [], + encrypted_content: 'wrapped:synthetic', + }); + expect(output[2]).toMatchObject({ event: { response: { output: [message, added.event.item] } } }); + }); + + test('does not synthesize affinity for a failed response', async () => { + const output: ProtocolFrame[] = []; + for await (const frame of wrapResponsesAffinityEgress(frames([ + eventFrame({ type: 'response.failed', response: response([], 'failed') }), + ]), { codec: immediateCodec, affinity })) output.push(frame); + + expect(output).toEqual([eventFrame({ type: 'response.failed', response: response([], 'failed') })]); + }); +}); From 2a9708c732ba9ffa2b9088eed6c7fde1e4b19255 Mon Sep 17 00:00:00 2001 From: Menci Date: Wed, 15 Jul 2026 19:52:55 +0800 Subject: [PATCH 007/157] feat(auth): persist per-key affinity secrets Generate a distinct hidden 32-byte secret for every API key and preserve it across key rotation. Backfill existing rows with a constrained D1/SQLite migration, expose the secret only through admin data transfer v10, and reject non-canonical imports. Keep normal key APIs secret-free and cover repository, migration, lifecycle, and transfer behavior. --- apps/platform-node/src/migrate_test.ts | 1 + .../src/components/settings/ExportSection.vue | 2 +- .../src/components/settings/ImportSection.vue | 2 +- .../0056_api_key_affinity_secret.sql | 43 +++++++++ packages/gateway/src/app-control_test.ts | 3 + .../src/control-plane/api-keys/routes.ts | 2 + .../src/control-plane/api-keys/routes_test.ts | 22 ++++- .../src/control-plane/data-transfer/routes.ts | 45 +++++++-- .../data-transfer/routes_test.ts | 91 ++++++++++++------- .../control-plane/performance/routes_test.ts | 6 ++ packages/gateway/src/control-plane/schemas.ts | 2 +- .../control-plane/search-usage/routes_test.ts | 1 + .../src/control-plane/telemetry-view_test.ts | 1 + .../control-plane/token-usage/routes_test.ts | 1 + .../gateway/src/control-plane/users/routes.ts | 2 + .../src/control-plane/users/routes_test.ts | 1 + .../chat/chat-completions/http_test.ts | 1 + .../src/data-plane/chat/gemini/http_test.ts | 1 + .../src/data-plane/chat/messages/http_test.ts | 1 + .../data-plane/chat/responses/http_test.ts | 1 + .../chat/shared/gateway-ctx_test.ts | 1 + packages/gateway/src/repo/api-keys_test.ts | 64 ++++++++++++- packages/gateway/src/repo/memory.ts | 9 +- packages/gateway/src/repo/sql.ts | 9 +- packages/gateway/src/repo/types.ts | 4 + .../gateway/src/shared/affinity-secret.ts | 11 +++ .../src/shared/affinity-secret_test.ts | 24 +++++ packages/gateway/src/test-helpers.ts | 1 + 28 files changed, 299 insertions(+), 53 deletions(-) create mode 100644 packages/gateway/migrations/0056_api_key_affinity_secret.sql create mode 100644 packages/gateway/src/shared/affinity-secret.ts create mode 100644 packages/gateway/src/shared/affinity-secret_test.ts diff --git a/apps/platform-node/src/migrate_test.ts b/apps/platform-node/src/migrate_test.ts index b3ef99c57..2cb97a980 100644 --- a/apps/platform-node/src/migrate_test.ts +++ b/apps/platform-node/src/migrate_test.ts @@ -32,6 +32,7 @@ test('applies all real migration files against a fresh sqlite', () => withTemp(a const colNames = apiKeyCols.results.map(r => r.name).toSorted(); assertEquals(colNames.includes('id'), true); assertEquals(colNames.includes('key'), true); + assertEquals(colNames.includes('affinity_secret'), true); // Every migration was recorded. const recorded = await db.prepare('SELECT COUNT(*) AS n FROM _migrations').first<{ n: number }>(); diff --git a/apps/web/src/components/settings/ExportSection.vue b/apps/web/src/components/settings/ExportSection.vue index 748110a48..94e9eb2f8 100644 --- a/apps/web/src/components/settings/ExportSection.vue +++ b/apps/web/src/components/settings/ExportSection.vue @@ -47,7 +47,7 @@ const exportData = async () => {