diff --git a/src/server/proxy-core/surfaces/sharedSurface.ts b/src/server/proxy-core/surfaces/sharedSurface.ts index 927b3b4..ec2c1c1 100644 --- a/src/server/proxy-core/surfaces/sharedSurface.ts +++ b/src/server/proxy-core/surfaces/sharedSurface.ts @@ -478,6 +478,7 @@ export function createSurfaceFailureToolkit(input: { promptTokens?: number | null; completionTokens?: number | null; totalTokens?: number | null; + usageSource?: 'upstream' | 'self-log' | 'unknown'; estimatedCost?: number; billingDetails?: unknown; upstreamPath?: string | null; @@ -497,6 +498,7 @@ export function createSurfaceFailureToolkit(input: { promptTokens: args.promptTokens, completionTokens: args.completionTokens, totalTokens: args.totalTokens, + usageSource: args.usageSource, estimatedCost: args.estimatedCost, billingDetails: args.billingDetails, upstreamPath: args.upstreamPath, diff --git a/src/server/proxy-core/surfaces/sharedSurface.usage-source.test.ts b/src/server/proxy-core/surfaces/sharedSurface.usage-source.test.ts new file mode 100644 index 0000000..2fecaba --- /dev/null +++ b/src/server/proxy-core/surfaces/sharedSurface.usage-source.test.ts @@ -0,0 +1,143 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const formatUtcSqlDateTimeMock = vi.fn(); +const composeProxyLogMessageMock = vi.fn(); +const insertProxyLogMock = vi.fn(); + +vi.mock('../../services/localTimeService.js', () => ({ + formatUtcSqlDateTime: (...args: unknown[]) => formatUtcSqlDateTimeMock(...args), +})); + +vi.mock('../../services/siteProxy.js', () => ({ + resolveChannelProxyUrl: vi.fn(), + withSiteRecordProxyRequestInit: vi.fn(), +})); + +vi.mock('../../services/tokenRouter.js', () => ({ + tokenRouter: { + recordFailure: vi.fn(), + recordSuccess: vi.fn(), + }, +})); + +vi.mock('../../services/proxyUsageFallbackService.js', () => ({ + resolveProxyUsageWithSelfLogFallback: vi.fn(), +})); + +vi.mock('../../services/alertService.js', () => ({ + reportProxyAllFailed: vi.fn(), + reportTokenExpired: vi.fn(), +})); + +vi.mock('../../services/alertRules.js', () => ({ + isTokenExpiredError: vi.fn(() => false), +})); + +vi.mock('../../services/proxyRetryPolicy.js', () => ({ + shouldRetryProxyRequest: vi.fn(() => false), +})); + +vi.mock('../../routes/proxy/logPathMeta.js', () => ({ + composeProxyLogMessage: (...args: unknown[]) => composeProxyLogMessageMock(...args), +})); + +vi.mock('../../routes/proxy/proxyBilling.js', () => ({ + resolveProxyLogBilling: vi.fn(), +})); + +vi.mock('../../services/proxyLogStore.js', () => ({ + insertProxyLog: (...args: unknown[]) => insertProxyLogMock(...args), +})); + +vi.mock('../../services/runtimeDispatch.js', () => ({ + dispatchRuntimeRequest: vi.fn(), +})); + +vi.mock('../orchestration/upstreamRequest.js', () => ({ + buildUpstreamUrl: vi.fn(), +})); + +vi.mock('../../services/oauth/quota.js', () => ({ + recordOauthQuotaResetHint: vi.fn(), +})); + +vi.mock('../../services/oauth/refreshSingleflight.js', () => ({ + refreshOauthAccessTokenSingleflight: vi.fn(), +})); + +vi.mock('../../services/proxyChannelCoordinator.js', () => ({ + proxyChannelCoordinator: { + buildStickySessionKey: vi.fn(), + getStickyChannelId: vi.fn(), + bindStickyChannel: vi.fn(), + clearStickyChannel: vi.fn(), + acquireChannelLease: vi.fn(), + }, +})); + +vi.mock('../executors/types.js', () => ({ + readRuntimeResponseText: vi.fn(), +})); + +vi.mock('../channelSelection.js', () => ({ + selectProxyChannelForAttempt: vi.fn(), +})); + +describe('shared surface usage source logging', () => { + beforeEach(() => { + formatUtcSqlDateTimeMock.mockReset(); + composeProxyLogMessageMock.mockReset(); + insertProxyLogMock.mockReset(); + }); + + it('forwards usage source through the failure toolkit log wrapper', async () => { + formatUtcSqlDateTimeMock.mockReturnValue('2026-04-05 21:00:00'); + composeProxyLogMessageMock.mockReturnValue('normalized error'); + insertProxyLogMock.mockResolvedValue(undefined); + + const { createSurfaceFailureToolkit } = await import('./sharedSurface.js'); + const toolkit = createSurfaceFailureToolkit({ + warningScope: 'responses', + downstreamPath: '/v1/responses', + maxRetries: 2, + clientContext: { + clientKind: 'codex', + sessionId: 'turn-1', + traceHint: 'trace-1', + }, + downstreamApiKeyId: null, + }); + + await toolkit.log({ + selected: { + channel: { id: 11, routeId: 22 }, + account: { id: 33, username: 'oauth-user' }, + site: { name: 'Codex OAuth' }, + actualModel: 'upstream-model', + }, + modelRequested: 'gpt-5.2', + status: 'success', + httpStatus: 200, + latencyMs: 320, + errorMessage: null, + retryCount: 0, + promptTokens: 12, + completionTokens: 4, + totalTokens: 16, + usageSource: 'self-log', + upstreamPath: '/v1/messages', + }); + + expect(composeProxyLogMessageMock).toHaveBeenCalledWith(expect.objectContaining({ + downstreamPath: '/v1/responses', + upstreamPath: '/v1/messages', + usageSource: 'self-log', + sessionId: 'turn-1', + traceHint: 'trace-1', + })); + expect(insertProxyLogMock).toHaveBeenCalledWith(expect.objectContaining({ + status: 'success', + errorMessage: 'normalized error', + })); + }); +}); diff --git a/src/server/routes/proxy/completions.ts b/src/server/routes/proxy/completions.ts index bec7795..1f227b6 100644 --- a/src/server/routes/proxy/completions.ts +++ b/src/server/routes/proxy/completions.ts @@ -208,6 +208,7 @@ export async function completionsProxyRoute(app: FastifyInstance) { billingDetails, clientContext, downstreamPath, + resolvedUsage.usageSource, isStream, firstByteLatencyMs, ); @@ -247,6 +248,7 @@ export async function completionsProxyRoute(app: FastifyInstance) { null, clientContext, downstreamPath, + null, isStream, firstByteLatencyMs, ); @@ -309,6 +311,7 @@ export async function completionsProxyRoute(app: FastifyInstance) { billingDetails, clientContext, downstreamPath, + resolvedUsage.usageSource, isStream, firstByteLatencyMs, ); @@ -338,6 +341,7 @@ export async function completionsProxyRoute(app: FastifyInstance) { null, clientContext, downstreamPath, + null, isStream, firstByteLatencyMs, ); @@ -381,6 +385,7 @@ async function logProxy( billingDetails: unknown = null, clientContext: DownstreamClientContext | null = null, downstreamPath = '/v1/completions', + usageSource: 'upstream' | 'self-log' | 'unknown' | null = null, isStream: boolean, firstByteLatencyMs: number | null, ) { @@ -393,6 +398,7 @@ async function logProxy( sessionId: clientContext?.sessionId || null, traceHint: clientContext?.traceHint || null, downstreamPath, + usageSource, errorMessage, }); await insertProxyLog({ diff --git a/src/server/routes/proxy/completions.usage-source.test.ts b/src/server/routes/proxy/completions.usage-source.test.ts new file mode 100644 index 0000000..9cb207d --- /dev/null +++ b/src/server/routes/proxy/completions.usage-source.test.ts @@ -0,0 +1,183 @@ +import Fastify, { type FastifyInstance } from 'fastify'; +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { mkdtempSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +const fetchMock = vi.fn(); +const selectChannelMock = vi.fn(); +const recordSuccessMock = vi.fn(); +const recordFailureMock = vi.fn(); +const refreshModelsAndRebuildRoutesMock = vi.fn(); +const reportProxyAllFailedMock = vi.fn(); +const reportTokenExpiredMock = vi.fn(); +const insertProxyLogMock = vi.fn(); +const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(); +const resolveProxyLogBillingMock = vi.fn(); + +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); + +vi.mock('../../services/tokenRouter.js', () => ({ + tokenRouter: { + selectChannel: (...args: unknown[]) => selectChannelMock(...args), + recordSuccess: (...args: unknown[]) => recordSuccessMock(...args), + recordFailure: (...args: unknown[]) => recordFailureMock(...args), + }, + invalidateTokenRouterCache: vi.fn(), +})); + +vi.mock('../../services/routeRefreshWorkflow.js', () => ({ + refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args), +})); + +vi.mock('../../services/alertService.js', () => ({ + reportProxyAllFailed: (...args: unknown[]) => reportProxyAllFailedMock(...args), + reportTokenExpired: (...args: unknown[]) => reportTokenExpiredMock(...args), +})); + +vi.mock('../../services/proxyLogStore.js', () => ({ + insertProxyLog: (...args: unknown[]) => insertProxyLogMock(...args), +})); + +vi.mock('../../services/proxyUsageFallbackService.js', () => ({ + resolveProxyUsageWithSelfLogFallback: (...args: unknown[]) => resolveProxyUsageWithSelfLogFallbackMock(...args), +})); + +vi.mock('./proxyBilling.js', () => ({ + resolveProxyLogBilling: (...args: unknown[]) => resolveProxyLogBillingMock(...args), +})); + +type DbModule = typeof import('../../db/index.js'); + +describe('/v1/completions usage source logging', () => { + let app: FastifyInstance; + let db: DbModule['db']; + let schema: DbModule['schema']; + let dataDir = ''; + + beforeAll(async () => { + dataDir = mkdtempSync(join(tmpdir(), 'metapi-completions-usage-source-')); + process.env.DATA_DIR = dataDir; + + await import('../../db/migrate.js'); + const dbModule = await import('../../db/index.js'); + const routesModule = await import('./completions.js'); + db = dbModule.db; + schema = dbModule.schema; + + app = Fastify(); + await app.register(routesModule.completionsProxyRoute); + }); + + beforeEach(async () => { + fetchMock.mockReset(); + selectChannelMock.mockReset(); + recordSuccessMock.mockReset(); + recordFailureMock.mockReset(); + refreshModelsAndRebuildRoutesMock.mockReset(); + reportProxyAllFailedMock.mockReset(); + reportTokenExpiredMock.mockReset(); + insertProxyLogMock.mockReset(); + resolveProxyUsageWithSelfLogFallbackMock.mockReset(); + resolveProxyLogBillingMock.mockReset(); + + resolveProxyUsageWithSelfLogFallbackMock.mockResolvedValue({ + promptTokens: 3, + completionTokens: 4, + totalTokens: 7, + usageSource: 'self-log', + }); + resolveProxyLogBillingMock.mockResolvedValue({ + estimatedCost: 0, + billingDetails: null, + }); + + await db.delete(schema.proxyLogs).run(); + await db.delete(schema.routeChannels).run(); + await db.delete(schema.tokenRoutes).run(); + await db.delete(schema.tokenModelAvailability).run(); + await db.delete(schema.modelAvailability).run(); + await db.delete(schema.accountTokens).run(); + await db.delete(schema.accounts).run(); + await db.delete(schema.siteApiEndpoints).run(); + await db.delete(schema.sites).run(); + }); + + afterAll(async () => { + await app.close(); + delete process.env.DATA_DIR; + }); + + it('stores usage source metadata on successful completion logs', async () => { + const site = await db.insert(schema.sites).values({ + name: 'usage-site', + url: 'https://console.example.com', + platform: 'new-api', + status: 'active', + }).returning().get(); + + const account = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'usage-user', + accessToken: '', + apiToken: 'sk-usage', + status: 'active', + checkinEnabled: false, + extraConfig: JSON.stringify({ credentialMode: 'apikey' }), + }).returning().get(); + + await db.insert(schema.siteApiEndpoints).values({ + siteId: site.id, + url: 'https://api-a.example.com', + enabled: true, + sortOrder: 0, + }).run(); + + selectChannelMock.mockResolvedValue({ + channel: { id: 11, routeId: 22 }, + site, + account, + tokenName: 'default', + tokenValue: 'sk-usage', + actualModel: 'gpt-4o-mini', + }); + + fetchMock.mockResolvedValue(new Response(JSON.stringify({ + id: 'cmpl-ok', + object: 'text_completion', + choices: [{ text: 'ok' }], + usage: { + prompt_tokens: 1, + completion_tokens: 2, + total_tokens: 3, + }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/completions', + headers: { + authorization: 'Bearer sk-downstream', + }, + payload: { + model: 'gpt-4o-mini', + prompt: 'hello', + }, + }); + + expect(response.statusCode).toBe(200); + expect(insertProxyLogMock).toHaveBeenCalledWith(expect.objectContaining({ + status: 'success', + errorMessage: expect.stringContaining('[usage:self-log]'), + })); + }); +}); diff --git a/src/server/routes/proxy/embeddings.siteApiEndpoint.test.ts b/src/server/routes/proxy/embeddings.siteApiEndpoint.test.ts new file mode 100644 index 0000000..4aeaaaa --- /dev/null +++ b/src/server/routes/proxy/embeddings.siteApiEndpoint.test.ts @@ -0,0 +1,199 @@ +import Fastify, { type FastifyInstance } from 'fastify'; +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { mkdtempSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +const fetchMock = vi.fn(); +const fetchWithObservedFirstByteMock = vi.fn(); +const getObservedResponseMetaMock = vi.fn(); +const selectChannelMock = vi.fn(); +const recordSuccessMock = vi.fn(); +const recordFailureMock = vi.fn(); +const refreshModelsAndRebuildRoutesMock = vi.fn(); +const reportProxyAllFailedMock = vi.fn(); +const reportTokenExpiredMock = vi.fn(); +const insertProxyLogMock = vi.fn(); +const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(); +const resolveProxyLogBillingMock = vi.fn(); + +vi.mock('undici', async () => { + const actual = await vi.importActual('undici'); + return { + ...actual, + fetch: (...args: unknown[]) => fetchMock(...args), + }; +}); + +vi.mock('../../proxy-core/firstByteTimeout.js', () => ({ + fetchWithObservedFirstByte: (...args: unknown[]) => fetchWithObservedFirstByteMock(...args), + getObservedResponseMeta: (...args: unknown[]) => getObservedResponseMetaMock(...args), +})); + +vi.mock('../../services/tokenRouter.js', () => ({ + tokenRouter: { + selectChannel: (...args: unknown[]) => selectChannelMock(...args), + recordSuccess: (...args: unknown[]) => recordSuccessMock(...args), + recordFailure: (...args: unknown[]) => recordFailureMock(...args), + }, + invalidateTokenRouterCache: vi.fn(), +})); + +vi.mock('../../services/routeRefreshWorkflow.js', () => ({ + refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args), +})); + +vi.mock('../../services/alertService.js', () => ({ + reportProxyAllFailed: (...args: unknown[]) => reportProxyAllFailedMock(...args), + reportTokenExpired: (...args: unknown[]) => reportTokenExpiredMock(...args), +})); + +vi.mock('../../services/proxyLogStore.js', () => ({ + insertProxyLog: (...args: unknown[]) => insertProxyLogMock(...args), +})); + +vi.mock('../../services/proxyUsageFallbackService.js', () => ({ + resolveProxyUsageWithSelfLogFallback: (...args: unknown[]) => resolveProxyUsageWithSelfLogFallbackMock(...args), +})); + +vi.mock('./proxyBilling.js', () => ({ + resolveProxyLogBilling: (...args: unknown[]) => resolveProxyLogBillingMock(...args), +})); + +type DbModule = typeof import('../../db/index.js'); + +describe('/v1/embeddings usage source logging', () => { + let app: FastifyInstance; + let db: DbModule['db']; + let schema: DbModule['schema']; + let dataDir = ''; + + beforeAll(async () => { + dataDir = mkdtempSync(join(tmpdir(), 'metapi-embeddings-site-api-endpoint-')); + process.env.DATA_DIR = dataDir; + + await import('../../db/migrate.js'); + const dbModule = await import('../../db/index.js'); + const routesModule = await import('./embeddings.js'); + db = dbModule.db; + schema = dbModule.schema; + + app = Fastify(); + await app.register(routesModule.embeddingsProxyRoute); + }); + + beforeEach(async () => { + fetchMock.mockReset(); + fetchWithObservedFirstByteMock.mockReset(); + getObservedResponseMetaMock.mockReset(); + selectChannelMock.mockReset(); + recordSuccessMock.mockReset(); + recordFailureMock.mockReset(); + refreshModelsAndRebuildRoutesMock.mockReset(); + reportProxyAllFailedMock.mockReset(); + reportTokenExpiredMock.mockReset(); + insertProxyLogMock.mockReset(); + resolveProxyUsageWithSelfLogFallbackMock.mockReset(); + resolveProxyLogBillingMock.mockReset(); + + fetchWithObservedFirstByteMock.mockImplementation(async (runner: (signal?: AbortSignal) => Promise) => runner()); + getObservedResponseMetaMock.mockReturnValue({ firstByteLatencyMs: 14 }); + resolveProxyUsageWithSelfLogFallbackMock.mockResolvedValue({ + promptTokens: 3, + completionTokens: 4, + totalTokens: 7, + usageSource: 'self-log', + }); + resolveProxyLogBillingMock.mockResolvedValue({ + estimatedCost: 0, + billingDetails: null, + }); + + await db.delete(schema.proxyLogs).run(); + await db.delete(schema.routeChannels).run(); + await db.delete(schema.tokenRoutes).run(); + await db.delete(schema.tokenModelAvailability).run(); + await db.delete(schema.modelAvailability).run(); + await db.delete(schema.accountTokens).run(); + await db.delete(schema.accounts).run(); + await db.delete(schema.siteApiEndpoints).run(); + await db.delete(schema.sites).run(); + }); + + afterAll(async () => { + await app.close(); + delete process.env.DATA_DIR; + }); + + it('stores usage source metadata on successful embedding logs', async () => { + const site = await db.insert(schema.sites).values({ + name: 'usage-site', + url: 'https://console.example.com', + platform: 'new-api', + status: 'active', + }).returning().get(); + + const account = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'usage-user', + accessToken: '', + apiToken: 'sk-usage', + status: 'active', + checkinEnabled: false, + extraConfig: JSON.stringify({ credentialMode: 'apikey' }), + }).returning().get(); + + await db.insert(schema.siteApiEndpoints).values({ + siteId: site.id, + url: 'https://api-a.example.com', + enabled: true, + sortOrder: 0, + }).run(); + + selectChannelMock.mockResolvedValue({ + channel: { id: 11, routeId: 22 }, + site, + account, + tokenName: 'default', + tokenValue: 'sk-usage', + actualModel: 'text-embedding-3-large', + }); + + fetchMock.mockResolvedValue(new Response(JSON.stringify({ + object: 'list', + data: [ + { + object: 'embedding', + embedding: [0.1, 0.2], + index: 0, + }, + ], + model: 'text-embedding-3-large', + usage: { + prompt_tokens: 1, + total_tokens: 1, + }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/embeddings', + headers: { + authorization: 'Bearer sk-downstream', + }, + payload: { + model: 'text-embedding-3-large', + input: 'hello', + }, + }); + + expect(response.statusCode).toBe(200); + expect(insertProxyLogMock).toHaveBeenCalledWith(expect.objectContaining({ + status: 'success', + errorMessage: expect.stringContaining('[usage:self-log]'), + })); + }); +}); diff --git a/src/server/routes/proxy/embeddings.ts b/src/server/routes/proxy/embeddings.ts index 35a361a..857fd91 100644 --- a/src/server/routes/proxy/embeddings.ts +++ b/src/server/routes/proxy/embeddings.ts @@ -155,7 +155,7 @@ export async function embeddingsProxyRoute(app: FastifyInstance) { logProxy( selected, requestedModel, 'success', upstream.status, latency, null, retryCount, downstreamApiKeyId, resolvedUsage.promptTokens, resolvedUsage.completionTokens, resolvedUsage.totalTokens, estimatedCost, billingDetails, clientContext, downstreamPath, - false, firstByteLatencyMs, + resolvedUsage.usageSource, false, firstByteLatencyMs, ); return reply.code(upstream.status).send(data); } catch (err: any) { @@ -183,6 +183,7 @@ export async function embeddingsProxyRoute(app: FastifyInstance) { null, clientContext, downstreamPath, + null, false, firstByteLatencyMs, ); @@ -229,6 +230,7 @@ async function logProxy( billingDetails: unknown = null, clientContext: DownstreamClientContext | null = null, downstreamPath = '/v1/embeddings', + usageSource: 'upstream' | 'self-log' | 'unknown' | null = null, isStream = false, firstByteLatencyMs: number | null = null, ) { @@ -239,10 +241,11 @@ async function logProxy( ? clientContext.clientKind : null, sessionId: clientContext?.sessionId || null, - traceHint: clientContext?.traceHint || null, - downstreamPath, - errorMessage, - }); + traceHint: clientContext?.traceHint || null, + downstreamPath, + usageSource, + errorMessage, + }); await insertProxyLog({ routeId: selected.channel.routeId, channelId: selected.channel.id, diff --git a/src/server/routes/proxy/endpointFlow.first-byte-timeout.test.ts b/src/server/routes/proxy/endpointFlow.first-byte-timeout.test.ts index 1192898..e17d568 100644 --- a/src/server/routes/proxy/endpointFlow.first-byte-timeout.test.ts +++ b/src/server/routes/proxy/endpointFlow.first-byte-timeout.test.ts @@ -15,15 +15,20 @@ function buildDelayedResponse( bodyText: string, delayMs: number, status = 200, - _signal?: AbortSignal, + signal?: AbortSignal, ): Response { const encoder = new TextEncoder(); const body = new ReadableStream({ start(controller) { - setTimeout(() => { + const timer = setTimeout(() => { + if (signal?.aborted) return; controller.enqueue(encoder.encode(bodyText)); controller.close(); }, delayMs); + + signal?.addEventListener('abort', () => { + clearTimeout(timer); + }, { once: true }); }, }); return new Response(body, {