Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[codex] fix empty stream success handling and unknown usage logs (#343)

Browse Source 
* fix: preserve repeated short responses deltas

* fix empty stream success handling and unknown usage logs

* fix proxy usage metadata typing

* fix stream bookkeeping and explicit zero usage source

* fix sticky fast path and usage placeholder detection
This commit is contained in:
Cita
2026-03-31 16:45:13 +08:00
committed by GitHub
parent 3ce5179fed
commit 6b19d06140
24 changed files with 926 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/proxy-core/surfaces/chatSurface.ts
+69 -24
View File
@@ -3,7 +3,7 @@ import type { FastifyRequest, FastifyReply } from 'fastify';
import { config } from '../../config.js';
import { tokenRouter } from '../../services/tokenRouter.js';
import { reportProxyAllFailed } from '../../services/alertService.js';
import { mergeProxyUsage, parseProxyUsage } from '../../services/proxyUsageParser.js';
import { hasProxyUsagePayload, mergeProxyUsage, parseProxyUsage } from '../../services/proxyUsageParser.js';
import { type DownstreamFormat } from '../../transformers/shared/normalized.js';
import {
buildClaudeCountTokensUpstreamRequest,
@@ -489,7 +489,10 @@ export async function handleChatSurfaceRequest(
if (isStream) {
const upstreamContentType = (upstream.headers.get('content-type') || '').toLowerCase();
let streamStarted = false;
const startSseResponse = () => {
if (streamStarted) return;
streamStarted = true;
reply.hijack();
reply.raw.statusCode = 200;
reply.raw.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
@@ -506,23 +509,54 @@ export async function handleChatSurfaceRequest(
cacheCreationTokens: 0,
promptTokensIncludeCache: null,
};
let upstreamUsagePresent = false;
const recordStreamSuccess = async (latencyMs: number) => {
await recordSurfaceSuccess({
selected,
requestedModel,
modelName,
parsedUsage,
upstreamUsagePresent,
requestStartedAtMs: startTime,
latencyMs,
retryCount,
upstreamPath: successfulUpstreamPath,
logSuccess: failureToolkit.log,
recordDownstreamCost: (estimatedCost) => {
recordDownstreamCostUsage(request, estimatedCost);
},
bestEffortMetrics: {
errorLabel: '[proxy/chat] failed to record success metrics',
},
});
};
const writeLines = (lines: string[]) => {
startSseResponse();
for (const line of lines) {
reply.raw.write(line);
}
};
const streamResponse = {
end() {
if (streamStarted) {
reply.raw.end();
}
},
};
const streamSession = openAiChatTransformer.proxyStream.createSession({
downstreamFormat,
modelName,
successfulUpstreamPath,
onParsedPayload: (payload) => {
if (payload && typeof payload === 'object') {
upstreamUsagePresent = upstreamUsagePresent || hasProxyUsagePayload(payload);
parsedUsage = mergeProxyUsage(parsedUsage, parseProxyUsage(payload));
}
},
writeLines,
writeRaw: (chunk) => {
startSseResponse();
reply.raw.write(chunk);
},
});
@@ -531,10 +565,9 @@ export async function handleChatSurfaceRequest(
const fallbackText = await readRuntimeResponseText(upstream);
rawText = fallbackText;
if (looksLikeResponsesSseText(fallbackText)) {
startSseResponse();
const streamResult = await streamSession.run(
createSingleChunkStreamReader(fallbackText),
reply.raw,
streamResponse,
);
const latency = Date.now() - startTime;
if (streamResult.status === 'failed') {
@@ -560,8 +593,17 @@ export async function handleChatSurfaceRequest(
type: 'stream_error',
},
}, successfulUpstreamPath);
if (!streamStarted) {
return reply.code(502).send({
error: {
message: streamResult.errorMessage,
type: 'upstream_error',
},
});
}
return;
}
await recordStreamSuccess(latency);
await finalizeDebugSuccess(
200,
successfulUpstreamPath,
@@ -588,6 +630,7 @@ export async function handleChatSurfaceRequest(
if (String(selected.site.platform || '').trim().toLowerCase() === 'gemini-cli') {
fallbackData = unwrapGeminiCliPayload(fallbackData);
}
upstreamUsagePresent = upstreamUsagePresent || hasProxyUsagePayload(fallbackData);
parsedUsage = mergeProxyUsage(parsedUsage, parseProxyUsage(fallbackData));
const latency = Date.now() - startTime;
const failure = detectProxyFailure({ rawText, usage: parsedUsage });
@@ -620,8 +663,7 @@ export async function handleChatSurfaceRequest(
return reply.code(failureOutcome.status).send(failureOutcome.payload);
}
startSseResponse();
const streamResult = streamSession.consumeUpstreamFinalPayload(fallbackData, fallbackText, reply.raw);
const streamResult = streamSession.consumeUpstreamFinalPayload(fallbackData, fallbackText, streamResponse);
if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
@@ -646,8 +688,17 @@ export async function handleChatSurfaceRequest(
type: 'stream_error',
},
}, successfulUpstreamPath);
if (!streamStarted) {
return reply.code(502).send({
error: {
message: streamResult.errorMessage,
type: 'upstream_error',
},
});
}
return;
}
await recordStreamSuccess(latency);
await finalizeDebugSuccess(
200,
successfulUpstreamPath,
@@ -663,8 +714,8 @@ export async function handleChatSurfaceRequest(
stickySessionKey,
selected,
});
return;
} else {
startSseResponse();
const upstreamReader = getRuntimeResponseReader(upstream);
const baseReader = String(selected.site.platform || '').trim().toLowerCase() === 'gemini-cli' && upstreamReader
? createGeminiCliStreamReader(upstreamReader)
@@ -687,7 +738,7 @@ export async function handleChatSurfaceRequest(
},
}
: baseReader;
const streamResult = await streamSession.run(reader, reply.raw);
const streamResult = await streamSession.run(reader, streamResponse);
rawText += decoder.decode();
const latency = Date.now() - startTime;
@@ -715,6 +766,14 @@ export async function handleChatSurfaceRequest(
type: 'stream_error',
},
}, successfulUpstreamPath);
if (!streamStarted) {
return reply.code(502).send({
error: {
message: streamResult.errorMessage,
type: 'upstream_error',
},
});
}
return;
}
@@ -725,23 +784,7 @@ export async function handleChatSurfaceRequest(
}
const latency = Date.now() - startTime;
await recordSurfaceSuccess({
selected,
requestedModel,
modelName,
parsedUsage,
requestStartedAtMs: startTime,
latencyMs: latency,
retryCount,
upstreamPath: successfulUpstreamPath,
logSuccess: failureToolkit.log,
recordDownstreamCost: (estimatedCost) => {
recordDownstreamCostUsage(request, estimatedCost);
},
bestEffortMetrics: {
errorLabel: '[proxy/chat] failed to record success metrics',
},
});
await recordStreamSuccess(latency);
await finalizeDebugSuccess(
200,
successfulUpstreamPath,
@@ -786,6 +829,7 @@ export async function handleChatSurfaceRequest(
const latency = Date.now() - startTime;
const parsedUsage = parseProxyUsage(upstreamData);
const upstreamUsagePresent = hasProxyUsagePayload(upstreamData);
const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) {
clearSurfaceStickyChannel({
@@ -823,6 +867,7 @@ export async function handleChatSurfaceRequest(
requestedModel,
modelName,
parsedUsage,
upstreamUsagePresent,
requestStartedAtMs: startTime,
latencyMs: latency,
retryCount,
src/server/proxy-core/surfaces/openAiResponsesSurface.ts
+23 -2
View File
@@ -2,7 +2,7 @@ import { TextDecoder } from 'node:util';
import type { FastifyReply, FastifyRequest } from 'fastify';
import { config } from '../../config.js';
import { reportProxyAllFailed } from '../../services/alertService.js';
import { mergeProxyUsage, parseProxyUsage } from '../../services/proxyUsageParser.js';
import { hasProxyUsagePayload, mergeProxyUsage, parseProxyUsage } from '../../services/proxyUsageParser.js';
import { openAiResponsesTransformer } from '../../transformers/openai/responses/index.js';
import {
buildUpstreamEndpointRequest,
@@ -604,13 +604,19 @@ export async function handleOpenAiResponsesSurfaceRequest(
const upstream = endpointResult.upstream;
const successfulUpstreamPath = endpointResult.upstreamPath;
const finalizeStreamSuccess = async (parsedUsage: UsageSummary, latency: number, streamDebugBody: unknown) => {
const finalizeStreamSuccess = async (
parsedUsage: UsageSummary,
latency: number,
streamDebugBody: unknown,
upstreamUsagePresent: boolean,
) => {
try {
await recordSurfaceSuccess({
selected,
requestedModel,
modelName,
parsedUsage,
upstreamUsagePresent,
requestStartedAtMs: startTime,
latencyMs: latency,
retryCount,
@@ -653,6 +659,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
cacheCreationTokens: 0,
promptTokensIncludeCache: null,
};
let upstreamUsagePresent = false;
const writeLines = (lines: string[]) => {
for (const line of lines) reply.raw.write(line);
};
@@ -663,6 +670,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
getUsage: () => parsedUsage,
onParsedPayload: (payload) => {
if (payload && typeof payload === 'object') {
upstreamUsagePresent = upstreamUsagePresent || hasProxyUsagePayload(payload);
parsedUsage = mergeProxyUsage(parsedUsage, parseProxyUsage(payload));
}
},
@@ -710,6 +718,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
parsedUsage,
latency,
debugTrace?.options.captureStreamChunks ? rawText : { stream: true, usage: parsedUsage },
upstreamUsagePresent,
);
bindSurfaceStickyChannel({
stickySessionKey,
@@ -728,6 +737,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
}
parsedUsage = parseProxyUsage(upstreamData);
upstreamUsagePresent = upstreamUsagePresent || hasProxyUsagePayload(upstreamData);
const latency = Date.now() - startTime;
const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) {
@@ -792,6 +802,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
parsedUsage,
latency,
debugTrace?.options.captureStreamChunks ? rawText : upstreamData,
upstreamUsagePresent,
);
bindSurfaceStickyChannel({
stickySessionKey,
@@ -808,6 +819,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
if (looksLikeResponsesSseText(rawText)) {
try {
const collectedPayload = collectResponsesFinalPayloadFromSseText(rawText, modelName).payload;
upstreamUsagePresent = upstreamUsagePresent || hasProxyUsagePayload(collectedPayload);
parsedUsage = mergeProxyUsage(parsedUsage, parseProxyUsage(collectedPayload));
const createdPayload = {
...collectedPayload,
@@ -829,7 +841,12 @@ export async function handleOpenAiResponsesSurfaceRequest(
parsedUsage,
latency,
debugTrace?.options.captureStreamChunks ? rawText : collectedPayload,
upstreamUsagePresent,
);
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return;
} catch {
// Fall through to the generic stream session for response.failed/error terminals.
@@ -867,6 +884,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
parsedUsage,
latency,
debugTrace?.options.captureStreamChunks ? rawText : { stream: true, usage: parsedUsage },
upstreamUsagePresent,
);
return;
}
@@ -937,6 +955,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
parsedUsage,
latency,
debugTrace?.options.captureStreamChunks ? rawText : { stream: true, usage: parsedUsage },
upstreamUsagePresent,
);
bindSurfaceStickyChannel({
stickySessionKey,
@@ -976,6 +995,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
}
const latency = Date.now() - startTime;
const parsedUsage = parseProxyUsage(upstreamData);
const upstreamUsagePresent = hasProxyUsagePayload(upstreamData);
const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) {
clearSurfaceStickyChannel({
@@ -1022,6 +1042,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
requestedModel,
modelName,
parsedUsage,
upstreamUsagePresent,
requestStartedAtMs: startTime,
latencyMs: latency,
retryCount,
src/server/proxy-core/surfaces/sharedSurface.test.ts
+68
View File
@@ -275,6 +275,7 @@ describe('selectSurfaceChannelForAttempt', () => {
estimatedCost: 0.42,
billingDetails: { source: 'test' },
upstreamPath: '/v1/responses',
usageSource: 'self-log',
clientContext: {
clientKind: 'codex',
clientAppId: 'app-id',
@@ -292,6 +293,7 @@ describe('selectSurfaceChannelForAttempt', () => {
traceHint: 'trace-1',
downstreamPath: '/v1/chat/completions',
upstreamPath: '/v1/responses',
usageSource: 'self-log',
errorMessage: 'upstream failed',
});
expect(insertProxyLogMock).toHaveBeenCalledWith({
@@ -828,6 +830,10 @@ describe('selectSurfaceChannelForAttempt', () => {
promptTokens: 20,
completionTokens: 8,
totalTokens: 28,
recoveredFromSelfLog: true,
estimatedCostFromQuota: 0.42,
selfLogBillingMeta: null,
usageSource: 'self-log',
});
resolveProxyLogBillingMock.mockResolvedValue({
estimatedCost: 0.42,
@@ -870,6 +876,7 @@ describe('selectSurfaceChannelForAttempt', () => {
requestStartedAtMs: 1000,
requestEndedAtMs: 1250,
localLatencyMs: 250,
upstreamUsagePresent: true,
usage: {
promptTokens: 10,
completionTokens: 5,
@@ -889,6 +896,10 @@ describe('selectSurfaceChannelForAttempt', () => {
promptTokens: 20,
completionTokens: 8,
totalTokens: 28,
recoveredFromSelfLog: true,
estimatedCostFromQuota: 0.42,
selfLogBillingMeta: null,
usageSource: 'self-log',
},
});
expect(recordSuccessMock).toHaveBeenCalledWith(11, 250, 0.42, 'upstream-model');
@@ -911,6 +922,7 @@ describe('selectSurfaceChannelForAttempt', () => {
promptTokens: 20,
completionTokens: 8,
totalTokens: 28,
usageSource: 'self-log',
estimatedCost: 0.42,
billingDetails: { source: 'pricing-test' },
upstreamPath: '/v1/responses',
@@ -920,12 +932,67 @@ describe('selectSurfaceChannelForAttempt', () => {
promptTokens: 20,
completionTokens: 8,
totalTokens: 28,
recoveredFromSelfLog: true,
estimatedCostFromQuota: 0.42,
selfLogBillingMeta: null,
usageSource: 'self-log',
},
estimatedCost: 0.42,
billingDetails: { source: 'pricing-test' },
});
});
it('logs unknown usage as null tokens while preserving success bookkeeping', async () => {
resolveProxyUsageWithSelfLogFallbackMock.mockResolvedValue({
promptTokens: 0,
completionTokens: 0,
totalTokens: 0,
recoveredFromSelfLog: false,
estimatedCostFromQuota: 0,
selfLogBillingMeta: null,
usageSource: 'unknown',
});
resolveProxyLogBillingMock.mockResolvedValue({
estimatedCost: 0,
billingDetails: null,
});
const logSuccess = vi.fn().mockResolvedValue(undefined);
const { recordSurfaceSuccess } = await import('./sharedSurface.js');
await recordSurfaceSuccess({
selected: {
channel: { id: 11, routeId: 22 },
account: { id: 33, username: 'oauth-user' },
site: { id: 44, url: 'https://upstream.example.com', platform: 'new-api', name: 'Upstream' },
tokenValue: 'live-token',
tokenName: 'default',
actualModel: 'upstream-model',
},
requestedModel: 'gpt-5.2',
modelName: 'upstream-model',
parsedUsage: {
promptTokens: 0,
completionTokens: 0,
totalTokens: 0,
},
requestStartedAtMs: 1000,
latencyMs: 250,
retryCount: 0,
upstreamPath: '/v1/chat/completions',
logSuccess,
});
expect(resolveProxyUsageWithSelfLogFallbackMock).toHaveBeenCalledWith(expect.objectContaining({
upstreamUsagePresent: false,
}));
expect(logSuccess).toHaveBeenCalledWith(expect.objectContaining({
promptTokens: null,
completionTokens: null,
totalTokens: null,
usageSource: 'unknown',
}));
});
it('treats success metrics as best-effort when requested', async () => {
resolveProxyUsageWithSelfLogFallbackMock.mockRejectedValueOnce(new Error('billing failed'));
const logSuccess = vi.fn().mockResolvedValue(undefined);
@@ -980,6 +1047,7 @@ describe('selectSurfaceChannelForAttempt', () => {
recoveredFromSelfLog: false,
estimatedCostFromQuota: 0,
selfLogBillingMeta: null,
usageSource: 'upstream',
},
estimatedCost: 0,
billingDetails: null,
src/server/proxy-core/surfaces/sharedSurface.ts
+45 -21
View File
@@ -98,6 +98,7 @@ type SurfaceResolvedUsageSummary = {
recoveredFromSelfLog: boolean;
estimatedCostFromQuota: number;
selfLogBillingMeta: import('../../services/proxyUsageFallbackService.js').SelfLogBillingMeta | null;
usageSource: 'upstream' | 'self-log' | 'unknown';
};
export async function selectSurfaceChannelForAttempt(input: {
@@ -228,12 +229,13 @@ export async function writeSurfaceProxyLog(input: {
errorMessage: string | null;
retryCount: number;
downstreamPath: string;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
promptTokens?: number | null;
completionTokens?: number | null;
totalTokens?: number | null;
estimatedCost?: number;
billingDetails?: unknown;
upstreamPath?: string | null;
usageSource?: 'upstream' | 'self-log' | 'unknown' | null;
clientContext?: DownstreamClientContext | null;
downstreamApiKeyId?: number | null;
}): Promise<void> {
@@ -247,6 +249,7 @@ export async function writeSurfaceProxyLog(input: {
traceHint: input.clientContext?.traceHint || null,
downstreamPath: input.downstreamPath,
upstreamPath: input.upstreamPath || null,
usageSource: input.usageSource || null,
errorMessage: input.errorMessage,
});
await insertProxyLog({
@@ -259,9 +262,9 @@ export async function writeSurfaceProxyLog(input: {
status: input.status,
httpStatus: input.httpStatus,
latencyMs: input.latencyMs,
promptTokens: input.promptTokens ?? 0,
completionTokens: input.completionTokens ?? 0,
totalTokens: input.totalTokens ?? 0,
promptTokens: input.promptTokens ?? null,
completionTokens: input.completionTokens ?? null,
totalTokens: input.totalTokens ?? null,
estimatedCost: input.estimatedCost ?? 0,
billingDetails: input.billingDetails ?? null,
clientFamily: input.clientContext?.clientKind || null,
@@ -354,6 +357,7 @@ export async function recordSurfaceSuccess(input: {
requestedModel: string;
modelName: string;
parsedUsage: SurfaceUsageSummary;
upstreamUsagePresent?: boolean;
requestStartedAtMs: number;
latencyMs: number;
retryCount: number;
@@ -366,9 +370,10 @@ export async function recordSurfaceSuccess(input: {
latencyMs: number;
errorMessage: string | null;
retryCount: number;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
promptTokens?: number | null;
completionTokens?: number | null;
totalTokens?: number | null;
usageSource?: 'upstream' | 'self-log' | 'unknown';
estimatedCost?: number;
billingDetails?: unknown;
upstreamPath?: string | null;
@@ -382,6 +387,11 @@ export async function recordSurfaceSuccess(input: {
estimatedCost: number;
billingDetails: unknown;
}> {
const hasUpstreamUsage = input.upstreamUsagePresent ?? (
input.parsedUsage.totalTokens > 0
|| input.parsedUsage.promptTokens > 0
|| input.parsedUsage.completionTokens > 0
);
let resolvedUsage: SurfaceResolvedUsageSummary = {
promptTokens: input.parsedUsage.promptTokens,
completionTokens: input.parsedUsage.completionTokens,
@@ -389,6 +399,7 @@ export async function recordSurfaceSuccess(input: {
recoveredFromSelfLog: false,
estimatedCostFromQuota: 0,
selfLogBillingMeta: null,
usageSource: hasUpstreamUsage ? 'upstream' : 'unknown',
};
let estimatedCost = 0;
let billingDetails: unknown = null;
@@ -403,6 +414,7 @@ export async function recordSurfaceSuccess(input: {
requestStartedAtMs: input.requestStartedAtMs,
requestEndedAtMs: input.requestStartedAtMs + input.latencyMs,
localLatencyMs: input.latencyMs,
upstreamUsagePresent: hasUpstreamUsage,
usage: {
promptTokens: input.parsedUsage.promptTokens,
completionTokens: input.parsedUsage.completionTokens,
@@ -432,6 +444,17 @@ export async function recordSurfaceSuccess(input: {
input.modelName,
);
input.recordDownstreamCost?.(estimatedCost);
const logTokens = resolvedUsage.usageSource === 'unknown'
? {
promptTokens: null,
completionTokens: null,
totalTokens: null,
}
: {
promptTokens: resolvedUsage.promptTokens,
completionTokens: resolvedUsage.completionTokens,
totalTokens: resolvedUsage.totalTokens,
};
await input.logSuccess({
selected: input.selected,
modelRequested: input.requestedModel,
@@ -440,9 +463,10 @@ export async function recordSurfaceSuccess(input: {
latencyMs: input.latencyMs,
errorMessage: null,
retryCount: input.retryCount,
promptTokens: resolvedUsage.promptTokens,
completionTokens: resolvedUsage.completionTokens,
totalTokens: resolvedUsage.totalTokens,
promptTokens: logTokens.promptTokens,
completionTokens: logTokens.completionTokens,
totalTokens: logTokens.totalTokens,
usageSource: resolvedUsage.usageSource,
estimatedCost,
billingDetails,
upstreamPath: input.upstreamPath,
@@ -470,9 +494,9 @@ export function createSurfaceFailureToolkit(input: {
latencyMs: number;
errorMessage: string | null;
retryCount: number;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
promptTokens?: number | null;
completionTokens?: number | null;
totalTokens?: number | null;
estimatedCost?: number;
billingDetails?: unknown;
upstreamPath?: string | null;
@@ -581,9 +605,9 @@ export function createSurfaceFailureToolkit(input: {
failure: { status: number; reason: string };
latencyMs: number;
retryCount: number;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
promptTokens?: number | null;
completionTokens?: number | null;
totalTokens?: number | null;
upstreamPath?: string | null;
}): Promise<SurfaceFailureOutcome> {
await tokenRouter.recordFailure(args.selected.channel.id, {
@@ -676,9 +700,9 @@ export function createSurfaceFailureToolkit(input: {
errorMessage: string | null;
latencyMs: number;
retryCount: number;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
promptTokens?: number | null;
completionTokens?: number | null;
totalTokens?: number | null;
upstreamPath?: string | null;
httpStatus?: number;
runtimeFailureStatus?: number | null;