diff --git a/.vscode/settings.json b/.vscode/settings.json index e997d2405..0a14f6b6e 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,8 +1,3 @@ { - "chatgpt.openOnStartup": true, - "i18n-ally.localesPaths": [ - "messages", - "src/i18n", - "src/app/[locale]/dashboard/sessions/[sessionId]/messages" - ] + "chatgpt.openOnStartup": true } diff --git a/src/app/v1/_lib/proxy/client-abort-metering.ts b/src/app/v1/_lib/proxy/client-abort-metering.ts index ea62e399b..937d4ae66 100644 --- a/src/app/v1/_lib/proxy/client-abort-metering.ts +++ b/src/app/v1/_lib/proxy/client-abort-metering.ts @@ -50,10 +50,12 @@ const USAGE_NUMBER_FIELDS = [ "thoughtsTokenCount", ] as const; +/** Narrows unknown JSON values to non-array records. */ function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } +/** Copies finite numeric accounting fields into a compact evidence record. */ function copyFiniteNumberFields( source: Record, target: Record, @@ -65,6 +67,7 @@ function copyFiniteNumberFields( } } +/** Compacts modality token details while discarding response content. */ function compactTokenDetails(value: unknown): unknown { if (!Array.isArray(value)) return undefined; const details = value.slice(0, 16).flatMap((entry) => { @@ -79,6 +82,7 @@ function compactTokenDetails(value: unknown): unknown { return details.length > 0 ? details : undefined; } +/** Retains only bounded usage and cache fields needed by billing. */ function compactUsage(value: unknown): Record | null { if (!isRecord(value)) return null; const compact: Record = {}; @@ -109,6 +113,7 @@ function compactUsage(value: unknown): Record | null { return Object.keys(compact).length > 0 ? compact : null; } +/** Retains a bounded protocol-error representation. */ function compactError(value: unknown): unknown { if (typeof value === "string") return value.slice(0, 1024); if (!isRecord(value)) return value === true ? true : undefined; @@ -120,6 +125,7 @@ function compactError(value: unknown): unknown { return Object.keys(compact).length > 0 ? compact : true; } +/** Builds a bounded frame payload without retaining generated content. */ function compactPayload(value: Record): Record { const compact: Record = {}; for (const field of [ @@ -185,6 +191,7 @@ function compactPayload(value: Record): Record return compact; } +/** Checks whether compact usage contains any positive billable count. */ function positiveUsage(value: unknown): boolean { const usage = compactUsage(value); if (!usage) return false; @@ -198,6 +205,7 @@ function positiveUsage(value: unknown): boolean { return false; } +/** Finds usage in supported provider envelopes. */ function findUsage(value: Record): boolean { if (positiveUsage(value.usage) || positiveUsage(value.usageMetadata)) return true; if (isRecord(value.message) && positiveUsage(value.message.usage)) return true; @@ -205,6 +213,7 @@ function findUsage(value: Record): boolean { return isRecord(value.response) && findUsage(value.response); } +/** Detects provider protocol-error markers in a parsed frame. */ function isProtocolError(value: Record): boolean { return ( value.error !== undefined || @@ -216,6 +225,7 @@ function isProtocolError(value: Record): boolean { ); } +/** Detects a terminal OpenAI Chat completion choice. */ function hasOpenAiCompletion(value: Record): boolean { return ( Array.isArray(value.choices) && @@ -228,6 +238,7 @@ function hasOpenAiCompletion(value: Record): boolean { ); } +/** Detects a terminal Gemini candidate finish reason. */ function hasGeminiCompletion(value: Record): boolean { const payload = isRecord(value.response) ? value.response : value; return ( @@ -241,6 +252,7 @@ function hasGeminiCompletion(value: Record): boolean { ); } +/** Incrementally frames SSE, data-only SSE, and bounded NDJSON input. */ class BoundedEventFramer { private readonly decoder = new TextDecoder("utf-8"); private line = ""; @@ -375,28 +387,45 @@ class BoundedEventFramer { } } +/** Creates the bounded accounting observer used after a client disconnect. */ export function createClientAbortMeteringObserver( format: ClientFormat ): ClientAbortMeteringObserver { const evidence = new Map(); + const evidenceBytes = new Map(); + const evidenceValueCounts = new Map(); const encoder = new TextEncoder(); + let retainedByteTotal = 0; let terminalSeen = false; let terminalUsageSeen = false; let protocolFailure: ClientAbortMeteringSnapshot["protocolFailure"] = null; let finished = false; - const retainedBytes = () => - [...new Set(evidence.values())].reduce( - (total, value) => total + encoder.encode(value).length, - 0 - ); - const setEvidence = (slot: EvidenceSlot, value: string): void => { const previous = evidence.get(slot); + if (previous === value) return; + const previousBytes = evidenceBytes.get(slot) ?? 0; + const nextBytes = encoder.encode(value).length; + const previousCount = previous ? (evidenceValueCounts.get(previous) ?? 0) : 0; + const nextCount = evidenceValueCounts.get(value) ?? 0; + const nextTotal = + retainedByteTotal - + (previousCount === 1 ? previousBytes : 0) + + (nextCount === 0 ? nextBytes : 0); + if (nextTotal > CLIENT_ABORT_METER_MAX_RETAINED_BYTES) return; + + if (previous) { + if (previousCount <= 1) { + evidenceValueCounts.delete(previous); + retainedByteTotal -= previousBytes; + } else { + evidenceValueCounts.set(previous, previousCount - 1); + } + } evidence.set(slot, value); - if (retainedBytes() <= CLIENT_ABORT_METER_MAX_RETAINED_BYTES) return; - evidence.delete(slot); - if (previous !== undefined) evidence.set(slot, previous); + evidenceBytes.set(slot, nextBytes); + evidenceValueCounts.set(value, nextCount + 1); + if (nextCount === 0) retainedByteTotal += nextBytes; }; const recordFrame = (frame: ParsedFrame): void => { @@ -508,7 +537,7 @@ export function createClientAbortMeteringObserver( return { text, billingComplete: isBillingComplete(), - retainedBytes: encoder.encode(text).length, + retainedBytes: retainedByteTotal, skippedOversizedFrames: framer.skippedOversizedFrames, protocolFailure: protocolFailure ? { ...protocolFailure } : null, }; diff --git a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts index 1f2a96ba7..f3d6a886d 100644 --- a/src/app/v1/_lib/proxy/demand-driven-response-pump.ts +++ b/src/app/v1/_lib/proxy/demand-driven-response-pump.ts @@ -159,6 +159,7 @@ export function createDemandDrivenResponsePump( settle(false, normalized, normalized); }; + /** Completes a detached drain after metering without reporting a source error. */ const finishDrain = (reason?: unknown) => { if (settled || state !== "draining") return; const normalized = reason == null ? new Error("Background drain complete") : toError(reason); diff --git a/src/app/v1/_lib/proxy/detached-stream-budget.test.ts b/src/app/v1/_lib/proxy/detached-stream-budget.test.ts index cfe38c290..0f31c454d 100644 --- a/src/app/v1/_lib/proxy/detached-stream-budget.test.ts +++ b/src/app/v1/_lib/proxy/detached-stream-budget.test.ts @@ -1,5 +1,8 @@ import { describe, expect, it } from "vitest"; -import { DetachedStreamBudget } from "./detached-stream-budget"; +import { + DetachedStreamBudget, + resolveDetachedStreamBudgetLimits, +} from "@/app/v1/_lib/proxy/detached-stream-budget"; function createBudget( overrides: Partial["limits"]> = {} @@ -70,4 +73,16 @@ describe("DetachedStreamBudget", () => { expect(() => budget.tryAcquire("metering", 0)).toThrow(RangeError); expect(budget.snapshot().activeStreams).toBe(0); }); + + it("uses conservative defaults when environment parsing fails", () => { + expect( + resolveDetachedStreamBudgetLimits(() => { + throw new Error("invalid environment"); + }) + ).toEqual({ + maxConcurrency: 64, + maxReservedBytes: 64 * 1024 * 1024, + meteringReserveBytes: 16 * 1024 * 1024, + }); + }); }); diff --git a/src/app/v1/_lib/proxy/detached-stream-budget.ts b/src/app/v1/_lib/proxy/detached-stream-budget.ts index 942100fde..791311c05 100644 --- a/src/app/v1/_lib/proxy/detached-stream-budget.ts +++ b/src/app/v1/_lib/proxy/detached-stream-budget.ts @@ -1,5 +1,9 @@ import { getEnvConfig } from "@/lib/config/env.schema"; +const DEFAULT_DETACHED_STREAM_MAX_CONCURRENCY = 64; +const DEFAULT_DETACHED_STREAM_BUDGET_BYTES = 64 * 1024 * 1024; +const DEFAULT_DETACHED_STREAM_METERING_RESERVE_BYTES = 16 * 1024 * 1024; + export type DetachedStreamLeaseKind = "metering" | "replay"; export interface DetachedStreamBudgetLimits { @@ -39,8 +43,10 @@ export class DetachedStreamBudget { private readonly activeByKind = createKindCounters(); private readonly reservedByKind = createKindCounters(); + /** Creates a process-local weighted detached-stream budget. */ constructor(private readonly resolveLimits: () => DetachedStreamBudgetLimits) {} + /** Attempts to reserve capacity for a metering or Replay detached stream. */ tryAcquire(kind: DetachedStreamLeaseKind, reservedBytes: number): DetachedStreamAcquireResult { if (!Number.isSafeInteger(reservedBytes) || reservedBytes <= 0) { throw new RangeError("Detached stream reservation must be a positive safe integer"); @@ -90,6 +96,7 @@ export class DetachedStreamBudget { }; } + /** Returns current usage and configured limits for diagnostics and tests. */ snapshot(): DetachedStreamBudgetSnapshot { return { activeStreams: this.activeStreams, @@ -103,21 +110,40 @@ export class DetachedStreamBudget { const DETACHED_STREAM_BUDGET_SYMBOL = Symbol.for("cch.detachedStreamBudget"); +/** Resolves configured detached-stream limits, falling back if env parsing fails. */ +export function resolveDetachedStreamBudgetLimits( + readEnv: () => ReturnType = getEnvConfig +): DetachedStreamBudgetLimits { + try { + const env = readEnv(); + return { + maxConcurrency: + env.DETACHED_STREAM_MAX_CONCURRENCY ?? DEFAULT_DETACHED_STREAM_MAX_CONCURRENCY, + maxReservedBytes: env.DETACHED_STREAM_BUDGET_BYTES ?? DEFAULT_DETACHED_STREAM_BUDGET_BYTES, + meteringReserveBytes: + env.DETACHED_STREAM_METERING_RESERVE_BYTES ?? + DEFAULT_DETACHED_STREAM_METERING_RESERVE_BYTES, + }; + } catch { + return { + maxConcurrency: DEFAULT_DETACHED_STREAM_MAX_CONCURRENCY, + maxReservedBytes: DEFAULT_DETACHED_STREAM_BUDGET_BYTES, + meteringReserveBytes: DEFAULT_DETACHED_STREAM_METERING_RESERVE_BYTES, + }; + } +} + function getDetachedStreamBudget(): DetachedStreamBudget { const globalState = globalThis as typeof globalThis & { [DETACHED_STREAM_BUDGET_SYMBOL]?: DetachedStreamBudget; }; - globalState[DETACHED_STREAM_BUDGET_SYMBOL] ??= new DetachedStreamBudget(() => { - const env = getEnvConfig(); - return { - maxConcurrency: env.DETACHED_STREAM_MAX_CONCURRENCY ?? 64, - maxReservedBytes: env.DETACHED_STREAM_BUDGET_BYTES ?? 64 * 1024 * 1024, - meteringReserveBytes: env.DETACHED_STREAM_METERING_RESERVE_BYTES ?? 16 * 1024 * 1024, - }; - }); + globalState[DETACHED_STREAM_BUDGET_SYMBOL] ??= new DetachedStreamBudget( + resolveDetachedStreamBudgetLimits + ); return globalState[DETACHED_STREAM_BUDGET_SYMBOL]; } +/** Acquires a weighted detached-stream lease from the process budget. */ export function acquireDetachedStreamLease( kind: DetachedStreamLeaseKind, reservedBytes: number @@ -125,6 +151,7 @@ export function acquireDetachedStreamLease( return getDetachedStreamBudget().tryAcquire(kind, reservedBytes); } +/** Returns the singleton detached-stream budget snapshot. */ export function getDetachedStreamBudgetSnapshot(): DetachedStreamBudgetSnapshot { return getDetachedStreamBudget().snapshot(); } diff --git a/src/app/v1/_lib/proxy/replay/replay-spool.ts b/src/app/v1/_lib/proxy/replay/replay-spool.ts index 12a54fa64..23318602f 100644 --- a/src/app/v1/_lib/proxy/replay/replay-spool.ts +++ b/src/app/v1/_lib/proxy/replay/replay-spool.ts @@ -44,7 +44,9 @@ function serializeDurablePersistence(operation: () => Promise): Promise } export interface ReplaySpoolOptions { + /** Shortens the detached window when the spool loses its active write role. */ onInactive?: () => void; + /** Runs once after completed, aborted, or disabled cleanup releases the spool. */ onTerminal?: () => void; } diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 973f7f7b5..33aa95cfd 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -106,8 +106,16 @@ function getSessionRequestOwnerKeyId(session: ProxySession): number | undefined return session.authState?.key?.id ?? session.messageContext?.key?.id ?? undefined; } -function resolveReplayDrainReservationBytes(): number { - const payloadBytes = getEnvConfig().REPLAY_MAX_PAYLOAD_BYTES; +/** Resolves a conservative Replay reservation for detached payload reconstruction. */ +export function resolveReplayDrainReservationBytes( + readEnv: () => ReturnType = getEnvConfig +): number { + let payloadBytes = 8 * 1024 * 1024; + try { + payloadBytes = readEnv().REPLAY_MAX_PAYLOAD_BYTES; + } catch { + // Keep a conservative reservation when environment parsing fails. + } // ReplaySpool keeps bounded write-back state, then terminal persistence reads // the Redis chunks and joins one payload string. Reserve the payload three // times for chunk strings, the joined string, and UTF-16 expansion. @@ -3851,41 +3859,48 @@ export class ProxyResponseHandler { AsyncTaskManager.touch(taskId); }; - const releaseTransportResources = () => { - if (transportReleased) return; - transportReleased = true; - clearPassthroughDrainTimeout(); - cleanupPassthroughClientAbortListener(); - cleanupTaskAbortBinding(); - clearIdleTimer(); - try { - const wasResponseControllerAborted = - sessionWithController.responseController?.signal.aborted ?? false; - const clientAborted = session.clientAbortSignal?.aborted ?? false; - const shouldClearTimeout = - responseTimeoutCleared || - streamEndedNormally || - wasResponseControllerAborted || - clientAborted; - if (shouldClearTimeout) { - clearResponseTimeoutOnce(); - } - } catch (error) { - logger.warn( - "[ResponseHandler] Gemini passthrough: Failed to clear response timeout", - { - taskId, - providerId: provider.id, - providerName: provider.name, - error: error instanceof Error ? error.message : String(error), + let transportReleasePromise: Promise | null = null; + const releaseTransportResources = async (): Promise => { + if (transportReleasePromise) return transportReleasePromise; + transportReleasePromise = (async () => { + await passthroughPump.teardown; + if (transportReleased) return; + transportReleased = true; + clearPassthroughDrainTimeout(); + cleanupPassthroughClientAbortListener(); + cleanupTaskAbortBinding(); + clearIdleTimer(); + try { + const wasResponseControllerAborted = + sessionWithController.responseController?.signal.aborted ?? false; + const clientAborted = session.clientAbortSignal?.aborted ?? false; + const shouldClearTimeout = + responseTimeoutCleared || + streamEndedNormally || + wasResponseControllerAborted || + clientAborted; + if (shouldClearTimeout) { + clearResponseTimeoutOnce(); } - ); - } - releaseSessionAgent(session); + } catch (error) { + logger.warn( + "[ResponseHandler] Gemini passthrough: Failed to clear response timeout", + { + taskId, + providerId: provider.id, + providerName: provider.name, + error: error instanceof Error ? error.message : String(error), + } + ); + } + releaseSessionAgent(session); + })(); + return transportReleasePromise; }; try { const pumpCompletion = await passthroughPump.completion; + await passthroughPump.teardown; streamEndedNormally = pumpCompletion.streamEndedNormally; pumpClientAborted = pumpCompletion.clientAborted; if (pumpCompletion.error) throw pumpCompletion.error; @@ -3895,7 +3910,7 @@ export class ProxyResponseHandler { const allContent = streamSnapshot.text; const clientAborted = pumpClientAborted || (session.clientAbortSignal?.aborted ?? false); - releaseTransportResources(); + await releaseTransportResources(); // 存储响应体到 Redis（5分钟过期） if ( @@ -3925,6 +3940,7 @@ export class ProxyResponseHandler { // 使用共享的统计处理方法 const duration = Date.now() - session.startTime; terminalFinalizationStarted = true; + const meteringSnapshot = passthroughClientDetached ? clientAbortMeter.finish() : null; const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( session, allContent, @@ -3933,12 +3949,12 @@ export class ProxyResponseHandler { clientAborted, discoveryLeaseLifecycle, streamProtocolObserver?.finish() ?? - (passthroughClientDetached + (meteringSnapshot ? { sawContent: false, - sawTerminal: clientAbortMeter.finish().billingComplete, - observationIncomplete: clientAbortMeter.finish().skippedOversizedFrames > 0, - failure: clientAbortMeter.finish().protocolFailure, + sawTerminal: meteringSnapshot.billingComplete, + observationIncomplete: meteringSnapshot.skippedOversizedFrames > 0, + failure: meteringSnapshot.protocolFailure, } : null), abortReason @@ -4005,6 +4021,7 @@ export class ProxyResponseHandler { clearIdleTimer(); const allContent = flushAndJoin(); const duration = Date.now() - session.startTime; + const meteringSnapshot = passthroughClientDetached ? clientAbortMeter.finish() : null; const finalized = await finalizeDeferredStreamingFinalizationIfNeeded( session, @@ -4013,12 +4030,12 @@ export class ProxyResponseHandler { false, clientAborted, discoveryLeaseLifecycle, - passthroughClientDetached + meteringSnapshot ? { sawContent: false, - sawTerminal: clientAbortMeter.finish().billingComplete, - observationIncomplete: clientAbortMeter.finish().skippedOversizedFrames > 0, - failure: clientAbortMeter.finish().protocolFailure, + sawTerminal: meteringSnapshot.billingComplete, + observationIncomplete: meteringSnapshot.skippedOversizedFrames > 0, + failure: meteringSnapshot.protocolFailure, } : null, abortReason @@ -4064,7 +4081,7 @@ export class ProxyResponseHandler { }); } } finally { - releaseTransportResources(); + await releaseTransportResources(); if (!commitSideEffectsScheduled) { void (async () => { await latestFinalizeAttemptResources?.(); @@ -4136,14 +4153,14 @@ export class ProxyResponseHandler { const decoder = new TextDecoder(); const text = decoder.decode(chunk, { stream: true }); buffer += text; - if (buffer.length > GEMINI_STREAM_TRANSFORM_MAX_BUFFER_CHARACTERS) { - buffer = ""; - throw new Error("Gemini stream line exceeded transform buffer limit"); - } const lines = buffer.split("\n"); // Keep the last line in buffer as it might be incomplete buffer = lines.pop() || ""; + if (buffer.length > GEMINI_STREAM_TRANSFORM_MAX_BUFFER_CHARACTERS) { + buffer = ""; + throw new Error("Gemini stream line exceeded transform buffer limit"); + } for (const line of lines) { const trimmedLine = line.trim(); @@ -4879,8 +4896,7 @@ export class ProxyResponseHandler { // 任何失败终态（假 200/中断/非 2xx）立即 abort，绝不被已完成重放命中。 if (replaySpool) { const activeReplaySpool = replaySpool; - const detachedReplayLease = - clientAbortDrainMode === "replay" ? clientAbortReplayLease : null; + const detachedReplayLease = clientAbortReplayLease; const isReplayableSuccess = finalized.commitSideEffects !== undefined && effectiveStatusCode >= 200 && @@ -4989,8 +5005,7 @@ export class ProxyResponseHandler { streamFinalizationPromise.catch(() => { if (replaySpool && !replaySpool.isTerminal) { streamReplayCompletionScheduled = true; - const detachedReplayLease = - clientAbortDrainMode === "replay" ? clientAbortReplayLease : null; + const detachedReplayLease = clientAbortReplayLease; void replaySpool .abort("finalize_error") .finally(() => releaseDetachedReplayLease(detachedReplayLease)); @@ -5110,6 +5125,7 @@ export class ProxyResponseHandler { const runProcessingTask = async () => { try { const pumpCompletion = await activeResponsePump.completion; + await activeResponsePump.teardown; cleanupTaskAbortBinding(); releaseSessionAgent(session); cleanupResponseControllerAbortListener(); @@ -5410,12 +5426,9 @@ export class ProxyResponseHandler { cleanupClientAbortListener(); clearClientAbortDrainTimer(); clearIdleTimer(); // 清除静默期计时器（防止泄漏） + await activeResponsePump.teardown; releaseSessionAgent(session); - if ( - clientAbortDrainMode === "replay" && - clientAbortReplayLease && - !streamReplayCompletionScheduled - ) { + if (clientAbortReplayLease && !streamReplayCompletionScheduled) { const detachedReplayLease = clientAbortReplayLease; if (replaySpool && !replaySpool.isTerminal) { void replaySpool diff --git a/src/lib/config/env.schema.ts b/src/lib/config/env.schema.ts index d2dbfdf96..04475acb4 100644 --- a/src/lib/config/env.schema.ts +++ b/src/lib/config/env.schema.ts @@ -29,6 +29,7 @@ const optionalNumber = (schema: z.ZodNumber) => /** * 环境变量验证schema */ +// biome-ignore format: preserve the established environment schema layout export const EnvSchema = z.object({ NODE_ENV: z.enum(["development", "production", "test"]).default("development"), DSN: optionalPreprocessed((val) => { @@ -197,7 +198,7 @@ export const EnvSchema = z.object({ DETACHED_STREAM_BUDGET_BYTES: z.coerce .number() .int() - .min(64 * 1024) + .min(3 * 1024 * 1024 + 64 * 1024) .max(1024 * 1024 * 1024) .default(64 * 1024 * 1024), DETACHED_STREAM_METERING_RESERVE_BYTES: z.coerce @@ -260,6 +261,14 @@ export const EnvSchema = z.object({ IP_GEO_API_TOKEN: z.string().optional(), IP_GEO_CACHE_TTL_SECONDS: z.coerce.number().int().min(60).max(86400).default(3600), IP_GEO_TIMEOUT_MS: z.coerce.number().int().min(100).max(10000).default(1500), +}).superRefine((env, context) => { + if (env.DETACHED_STREAM_METERING_RESERVE_BYTES > env.DETACHED_STREAM_BUDGET_BYTES) { + context.addIssue({ + code: "custom", + path: ["DETACHED_STREAM_METERING_RESERVE_BYTES"], + message: "DETACHED_STREAM_METERING_RESERVE_BYTES cannot exceed DETACHED_STREAM_BUDGET_BYTES", + }); + } }); /** diff --git a/tests/unit/lib/env-detached-stream-budget.test.ts b/tests/unit/lib/env-detached-stream-budget.test.ts index deba36098..7ddb70ecb 100644 --- a/tests/unit/lib/env-detached-stream-budget.test.ts +++ b/tests/unit/lib/env-detached-stream-budget.test.ts @@ -12,17 +12,38 @@ describe("EnvSchema - detached stream budget", () => { it("parses explicit budget limits", () => { const env = EnvSchema.parse({ DETACHED_STREAM_MAX_CONCURRENCY: "8", - DETACHED_STREAM_BUDGET_BYTES: String(512 * 1024), + DETACHED_STREAM_BUDGET_BYTES: String(4 * 1024 * 1024), DETACHED_STREAM_METERING_RESERVE_BYTES: String(128 * 1024), }); expect(env.DETACHED_STREAM_MAX_CONCURRENCY).toBe(8); - expect(env.DETACHED_STREAM_BUDGET_BYTES).toBe(512 * 1024); + expect(env.DETACHED_STREAM_BUDGET_BYTES).toBe(4 * 1024 * 1024); expect(env.DETACHED_STREAM_METERING_RESERVE_BYTES).toBe(128 * 1024); }); it("rejects a budget smaller than one metering reservation", () => { expect(() => - EnvSchema.parse({ DETACHED_STREAM_BUDGET_BYTES: String(64 * 1024 - 1) }) + EnvSchema.parse({ + DETACHED_STREAM_BUDGET_BYTES: String(3 * 1024 * 1024 + 64 * 1024 - 1), + }) ).toThrow(); }); + + it("rejects a metering reserve larger than the total budget", () => { + expect(() => + EnvSchema.parse({ + DETACHED_STREAM_BUDGET_BYTES: String(4 * 1024 * 1024), + DETACHED_STREAM_METERING_RESERVE_BYTES: String(5 * 1024 * 1024), + }) + ).toThrow(); + }); + + it("allows a metering reserve equal to the total budget", () => { + const budget = 4 * 1024 * 1024; + expect( + EnvSchema.parse({ + DETACHED_STREAM_BUDGET_BYTES: String(budget), + DETACHED_STREAM_METERING_RESERVE_BYTES: String(budget), + }).DETACHED_STREAM_METERING_RESERVE_BYTES + ).toBe(budget); + }); }); diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index 92f32b2e1..cc9563e74 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -7,6 +7,7 @@ import { import { BoundedStreamTextAccumulator, ProxyResponseHandler, + resolveReplayDrainReservationBytes, } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; import { @@ -1109,6 +1110,14 @@ describe("ProxyResponseHandler stream client abort finalization", () => { ); }); + it("uses a conservative Replay reservation when environment parsing fails", () => { + expect( + resolveReplayDrainReservationBytes(() => { + throw new Error("invalid environment"); + }) + ).toBe(29 * 1024 * 1024); + }); + it("propagates unexpected registered task rejections during drain", async () => { const failure = new Error("factory task failed"); AsyncTaskManager.register("rejecting-test-task", async () => { diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts index 76d8c5b14..8532dc255 100644 --- a/tests/unit/proxy/response-handler-stream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -615,6 +615,25 @@ describe("ProxyResponseHandler.dispatch stream terminal behavior", () => { ); }); + it("accepts a large Gemini chunk composed of complete short lines", async () => { + const { session } = await createSession({}); + session.setProvider({ ...createProvider(), providerType: "gemini" }); + session.originalFormat = "claude"; + const frame = 'data: {"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}\n\n'; + const body = `${frame.repeat(Math.ceil((1024 * 1024 + 1) / frame.length))}data: {"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}\n\n`; + + const returned = await ProxyResponseHandler.dispatch(session, sseResponse(body)); + const returnedText = await returned.text(); + expect(returnedText.length).toBeGreaterThan(1024 * 1024); + await settleTasks(); + + expect(mocks.durable).toHaveBeenCalledWith( + MESSAGE.id, + expect.objectContaining({ statusCode: 200 }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + it("treats malformed after wrapped Gemini content as postcommit", async () => { const { session } = await createSession({}); session.setProvider({ ...createProvider(), providerType: "gemini" });