import { performance } from "node:perf_hooks";import { Effect, Schema, Stream } from "effect";import type { AppContext } from "../../app-context";import { buildSseHeaders } from "../../http/sse";import type { ProviderRouteConfig } from "../../services/provider-routing";import type { Recipe } from "../models/types";import { getDefaultReasoningParser } from "../compute/recipe-defaults";import { shouldBufferImplicitReasoningContent } from "./reasoning";import { recordStreamingInferenceUsage } from "./inference-accounting";import { createToolCallStream, type StreamUsage } from "./tool-call-stream"; const KEEPALIVE_INTERVAL_MS = 15_000; export class ChatCompletionsStreamError extends Schema.TaggedErrorClass<ChatCompletionsStreamError>()( "ChatCompletionsStreamError", { stage: Schema.Literals(["connect", "response", "stream"]), message: Schema.String, source: Schema.optional(Schema.Unknown), },) {} export interface ChatCompletionsStreamParameters { upstreamUrl: string; headers: Record<string, string>; body: BodyInit; clientSignal: AbortSignal; matchedRecipe: Recipe | null; sourceHeader: string | null; sessionId: string | null; recordedModel: string; recordedProvider: string; requestStart: number; requestProvider: string; providerRouting: ProviderRouteConfig | null; context: Pick<AppContext, "logger" | "stores">; keepaliveIntervalMs?: number;} const errorFrame = (message: string): Uint8Array => new TextEncoder().encode( `data: ${JSON.stringify({ error: { message, type: "upstream_error" } })}\n\n`, ); const responseErrorFrame = (status: number, body: string): Uint8Array => new TextEncoder().encode( `data: ${body || JSON.stringify({ error: { message: `Upstream returned ${status}`, type: "upstream_error" } })}\n\n`, ); export const shouldBufferImplicitReasoning = (input: { matchedRecipe: Recipe | null; recordedModel: string;}): boolean => { const { matchedRecipe } = input; const reasoningParser = matchedRecipe && matchedRecipe.reasoning_parser !== null ? matchedRecipe.reasoning_parser : matchedRecipe ? (getDefaultReasoningParser(matchedRecipe) ?? null) : null; const upstreamParsesReasoning = (matchedRecipe?.backend === "vllm" || matchedRecipe?.backend === "sglang") && Boolean(reasoningParser); return ( !upstreamParsesReasoning && shouldBufferImplicitReasoningContent(input.recordedModel, reasoningParser) );}; const responseBodyStream = ( upstreamResponse: Response, parameters: ChatCompletionsStreamParameters,): Stream.Stream<Uint8Array, never> => { const { matchedRecipe, sourceHeader, sessionId, recordedModel, recordedProvider, requestStart, providerRouting, requestProvider, context, } = parameters; const source = upstreamResponse.body; if (!source) { return Stream.succeed( errorFrame( providerRouting ? `${requestProvider} backend unavailable` : "Inference backend unavailable", ), ); } let ttftMs: number | null = null; let observedUsage: StreamUsage | null = null; const transformed = createToolCallStream( source, (usage) => { observedUsage = usage; }, () => { ttftMs ??= Math.max(0, Math.round(performance.now() - requestStart)); }, { bufferImplicitReasoningContent: shouldBufferImplicitReasoning({ matchedRecipe, recordedModel, }), }, ); return Stream.fromReadableStream({ evaluate: () => transformed, onError: (source) => new ChatCompletionsStreamError({ stage: "stream", message: "Chat completions stream failed", source, }), }).pipe( Stream.catchCause((cause) => { if (!parameters.clientSignal.aborted) { context.logger.error("Stream pipe error", { error: String(cause) }); } return Stream.empty; }), Stream.ensuring( Effect.suspend(() => observedUsage ? recordStreamingInferenceUsage( { logger: context.logger, stores: context.stores }, { usage: observedUsage, record: { model: recordedModel, source: sourceHeader, session_id: sessionId, provider: recordedProvider, ttft_ms: ttftMs, duration_ms: Math.round(performance.now() - requestStart), status: upstreamResponse.status, }, }, ).pipe( Effect.catch((error) => Effect.sync(() => context.logger.warn("Streaming accounting failed", { error: String(error) }), ), ), ) : Effect.void, ), ), );}; const upstreamStream = ( parameters: ChatCompletionsStreamParameters,): Stream.Stream<Uint8Array, never> => Stream.unwrap( Effect.tryPromise({ try: (signal) => fetch(parameters.upstreamUrl, { method: "POST", headers: parameters.headers, body: parameters.body, signal: AbortSignal.any([parameters.clientSignal, signal]), }), catch: (source) => new ChatCompletionsStreamError({ stage: "connect", message: "Chat completions connection failed", source, }), }).pipe( Effect.flatMap((response) => { if (response.ok) return Effect.succeed(responseBodyStream(response, parameters)); return Effect.tryPromise({ try: () => response.text(), catch: (source) => new ChatCompletionsStreamError({ stage: "response", message: "Chat completions response failed", source, }), }).pipe( Effect.map((body) => Stream.succeed(responseErrorFrame(response.status, body))), Effect.catch(() => Effect.succeed(Stream.succeed(responseErrorFrame(response.status, ""))), ), ); }), Effect.catch((error) => Effect.succeed( parameters.clientSignal.aborted ? Stream.empty : Stream.succeed(errorFrame(`Upstream connection failed: ${error.message}`)), ), ), ), ); export const buildChatCompletionsStreamResponse = ( parameters: ChatCompletionsStreamParameters,): Response => { const keepalive = new TextEncoder().encode(": keepalive\n\n"); const heartbeat = Stream.concat( Stream.succeed(keepalive), Stream.tick(parameters.keepaliveIntervalMs ?? KEEPALIVE_INTERVAL_MS).pipe( Stream.map(() => keepalive), ), ); const stream = Stream.merge(upstreamStream(parameters), heartbeat, { haltStrategy: "left" }); return new Response(Stream.toReadableStream(stream), { headers: buildSseHeaders() });};
import { performance } from "node:perf_hooks";
import { Effect, Schema, Stream } from "effect";
import type { AppContext } from "../../app-context";
import { buildSseHeaders } from "../../http/sse";
import type { ProviderRouteConfig } from "../../services/provider-routing";
import type { Recipe } from "../models/types";
import { getDefaultReasoningParser } from "../compute/recipe-defaults";
import { shouldBufferImplicitReasoningContent } from "./reasoning";
import { recordStreamingInferenceUsage } from "./inference-accounting";
import { createToolCallStream, type StreamUsage } from "./tool-call-stream";
const KEEPALIVE_INTERVAL_MS = 15_000;
export class ChatCompletionsStreamError extends Schema.TaggedErrorClass<ChatCompletionsStreamError>()(
"ChatCompletionsStreamError",
{
stage: Schema.Literals(["connect", "response", "stream"]),
message: Schema.String,
source: Schema.optional(Schema.Unknown),
},
) {}
export interface ChatCompletionsStreamParameters {
upstreamUrl: string;
headers: Record<string, string>;
body: BodyInit;
clientSignal: AbortSignal;
matchedRecipe: Recipe | null;
sourceHeader: string | null;
sessionId: string | null;
recordedModel: string;
recordedProvider: string;
requestStart: number;
requestProvider: string;
providerRouting: ProviderRouteConfig | null;
context: Pick<AppContext, "logger" | "stores">;
keepaliveIntervalMs?: number;
}
const errorFrame = (message: string): Uint8Array =>
new TextEncoder().encode(
`data: ${JSON.stringify({ error: { message, type: "upstream_error" } })}\n\n`,
);
const responseErrorFrame = (status: number, body: string): Uint8Array =>
new TextEncoder().encode(
`data: ${body || JSON.stringify({ error: { message: `Upstream returned ${status}`, type: "upstream_error" } })}\n\n`,
);
export const shouldBufferImplicitReasoning = (input: {
matchedRecipe: Recipe | null;
recordedModel: string;
}): boolean => {
const { matchedRecipe } = input;
const reasoningParser =
matchedRecipe && matchedRecipe.reasoning_parser !== null
? matchedRecipe.reasoning_parser
: matchedRecipe
? (getDefaultReasoningParser(matchedRecipe) ?? null)
: null;
const upstreamParsesReasoning =
(matchedRecipe?.backend === "vllm" || matchedRecipe?.backend === "sglang") &&
Boolean(reasoningParser);
return (
!upstreamParsesReasoning &&
shouldBufferImplicitReasoningContent(input.recordedModel, reasoningParser)
);
};
const responseBodyStream = (
upstreamResponse: Response,
parameters: ChatCompletionsStreamParameters,
): Stream.Stream<Uint8Array, never> => {
const {
matchedRecipe,
sourceHeader,
sessionId,
recordedModel,
recordedProvider,
requestStart,
providerRouting,
requestProvider,
context,
} = parameters;
const source = upstreamResponse.body;
if (!source) {
return Stream.succeed(
errorFrame(
providerRouting
? `${requestProvider} backend unavailable`
: "Inference backend unavailable",
),
);
}
let ttftMs: number | null = null;
let observedUsage: StreamUsage | null = null;
const transformed = createToolCallStream(
source,
(usage) => {
observedUsage = usage;
},
() => {
ttftMs ??= Math.max(0, Math.round(performance.now() - requestStart));
},
{
bufferImplicitReasoningContent: shouldBufferImplicitReasoning({
matchedRecipe,
recordedModel,
}),
},
);
return Stream.fromReadableStream({
evaluate: () => transformed,
onError: (source) =>
new ChatCompletionsStreamError({
stage: "stream",
message: "Chat completions stream failed",
source,
}),
}).pipe(
Stream.catchCause((cause) => {
if (!parameters.clientSignal.aborted) {
context.logger.error("Stream pipe error", { error: String(cause) });
}
return Stream.empty;
}),
Stream.ensuring(
Effect.suspend(() =>
observedUsage
? recordStreamingInferenceUsage(
{ logger: context.logger, stores: context.stores },
{
usage: observedUsage,
record: {
model: recordedModel,
source: sourceHeader,
session_id: sessionId,
provider: recordedProvider,
ttft_ms: ttftMs,
duration_ms: Math.round(performance.now() - requestStart),
status: upstreamResponse.status,
},
},
).pipe(
Effect.catch((error) =>
Effect.sync(() =>
context.logger.warn("Streaming accounting failed", { error: String(error) }),
),
),
)
: Effect.void,
),
),
);
};
const upstreamStream = (
parameters: ChatCompletionsStreamParameters,
): Stream.Stream<Uint8Array, never> =>
Stream.unwrap(
Effect.tryPromise({
try: (signal) =>
fetch(parameters.upstreamUrl, {
method: "POST",
headers: parameters.headers,
body: parameters.body,
signal: AbortSignal.any([parameters.clientSignal, signal]),
}),
catch: (source) =>
new ChatCompletionsStreamError({
stage: "connect",
message: "Chat completions connection failed",
source,
}),
}).pipe(
Effect.flatMap((response) => {
if (response.ok) return Effect.succeed(responseBodyStream(response, parameters));
return Effect.tryPromise({
try: () => response.text(),
catch: (source) =>
new ChatCompletionsStreamError({
stage: "response",
message: "Chat completions response failed",
source,
}),
}).pipe(
Effect.map((body) => Stream.succeed(responseErrorFrame(response.status, body))),
Effect.catch(() =>
Effect.succeed(Stream.succeed(responseErrorFrame(response.status, ""))),
),
);
}),
Effect.catch((error) =>
Effect.succeed(
parameters.clientSignal.aborted
? Stream.empty
: Stream.succeed(errorFrame(`Upstream connection failed: ${error.message}`)),
),
),
),
);
export const buildChatCompletionsStreamResponse = (
parameters: ChatCompletionsStreamParameters,
): Response => {
const keepalive = new TextEncoder().encode(": keepalive\n\n");
const heartbeat = Stream.concat(
Stream.succeed(keepalive),
Stream.tick(parameters.keepaliveIntervalMs ?? KEEPALIVE_INTERVAL_MS).pipe(
Stream.map(() => keepalive),
),
);
const stream = Stream.merge(upstreamStream(parameters), heartbeat, { haltStrategy: "left" });
return new Response(Stream.toReadableStream(stream), { headers: buildSseHeaders() });
};