dcosnet
/
SysDeck
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
SysDeck/klanker-gate/apps/gateway/routes/advanced.ts

695 lines
24 KiB
TypeScript
Executable File
Raw Blame History

import {
EmbeddingRequestSchema,
type ImageGenerationRequest,
ImageGenerationRequestSchema,
MAX_IMAGE_N,
MAX_MEDIA_INFLIGHT_BYTES,
MAX_TRANSCRIPTION_JSON_BYTES,
MEDIA_BYTES_PER_IMAGE,
type ProviderCapability,
SpeechRequestSchema,
} from "../../../packages/contracts/src/mod.ts";
import {
errorResponse,
GatewayError,
type Router,
} from "../../../packages/core/src/mod.ts";
import {
ProviderError,
type ResolvedTarget,
targetScope,
} from "../../../packages/providers/src/mod.ts";
import {
clampCount,
clampSeconds,
type CounterSink,
type MediaUnits,
mergeRequestStatus,
mergeRequestTokens,
mergeRequestUnits,
normalizeUsage,
type RequestDispatch,
setRequestDispatch,
type UsageShape,
} from "../../../packages/telemetry/src/usage.ts";
import type { AppContext } from "../context.ts";
import {
dispatchDetached,
dispatchThenDetach,
jsonResponse,
mapDispatchError,
parseJsonBody,
readCappedBytes,
rebuild,
validationErrorResponse,
} from "./helpers.ts";
/** The three media surfaces. Used both to register the routes and to gate the
* unit readers, so a path change cannot desync the gate from the route. */
const IMAGES_PATH = "/v1/images/generations";
const SPEECH_PATH = "/v1/audio/speech";
const TRANSCRIPTIONS_PATH = "/v1/audio/transcriptions";
function requireCapability(
target: ResolvedTarget,
capability: keyof ProviderCapability,
what: string,
): void {
if (!target.capabilities[capability]) {
throw new GatewayError(
400,
`Provider "${target.providerId}" does not support ${what}.`,
"invalid_request_error",
);
}
}
/** Routes by provider prefix in the model, else default-provider rules. */
function targetFromModel(ctx: AppContext, model: string): ResolvedTarget {
return ctx.providers.resolve(model);
}
/** Routes multipart/list endpoints by ?provider= or the default provider. */
function targetFromQuery(ctx: AppContext, req: Request): ResolvedTarget {
const provider = new URL(req.url).searchParams.get("provider");
const target = provider
? ctx.providers.accountTarget(provider)
: ctx.providers.resolve("");
// A provider-scoped virtual key must not read files/batches from an
// out-of-scope account via ?provider= (these data-plane paths carry no model,
// so only the provider allowlist applies).
const scope = targetScope.getStore();
if (scope?.providers && !scope.providers.has(target.providerId)) {
throw new GatewayError(
403,
`Virtual key is not permitted to use provider "${target.providerId}".`,
"governance_error",
undefined,
"provider_not_permitted",
);
}
return target;
}
function jsonForward(path: string, body: unknown): Request {
return new Request(`http://internal${path}`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
});
}
function requireRawProxy(target: ResolvedTarget) {
if (!target.adapter.rawProxy) {
throw new GatewayError(
400,
`Provider "${target.providerId}" has no passthrough surface for this endpoint.`,
"invalid_request_error",
);
}
return target.adapter.rawProxy.bind(target.adapter);
}
// ---------------------------------------------------------------------------
// Media accounting: the in-flight byte budget, the quantity readers, and the
// one dispatcher that writes the channel.
/**
* Bytes reserved across every media body read in flight in THIS process.
* Reserved pessimistically at the per-response cap before the provider is
* dispatched to, and held until the read finishes, so a mid-read crossing is
* structurally impossible. Because the hold spans the provider's generation
* latency it is also the media path's only concurrency bound.
*/
let mediaInflightBytes = 0;
/**
* Bytes currently reserved. Nothing in production reads this; it exists so the
* `finally` release is observable, since a reservation that leaks on one exit
* path degrades into a permanent 429 and nothing else would show it.
*/
export function mediaInflightReserved(): number {
return mediaInflightBytes;
}
type MediaReservation =
| { ok: true; bytes: number }
| { ok: false; response: Response };
/**
* Bytes ONE response on this surface may be buffered into, which is also the
* reservation it takes. Zero means the surface is not buffered at all: a TTS
* reply streams straight through, so it needs no cap and takes no reservation.
*
* `sampleCount` before `n`, matching `imagen.ts`. `Math.min` is defence in
* depth, not the bound - the bound is the schema's `.max(MAX_IMAGE_N)`, which
* every caller of this function has already applied.
*/
function mediaReservationBytes(
pathname: string,
body?: ImageGenerationRequest,
): number {
if (pathname === IMAGES_PATH) {
const requested = clampCount(body?.sampleCount ?? body?.n ?? 1) || 1;
return Math.min(requested, MAX_IMAGE_N) * MEDIA_BYTES_PER_IMAGE;
}
if (pathname === TRANSCRIPTIONS_PATH) {
return MAX_TRANSCRIPTION_JSON_BYTES;
}
return 0;
}
/**
* Takes `bytes` from the process budget, or refuses with the house 429. The
* refusal happens before the provider is reached, so nothing was spent and
* there is nothing to bill. Fail-closed: over budget denies.
*/
function reserveMediaInflight(
bytes: number,
counters: CounterSink,
): MediaReservation {
if (bytes <= 0) {
return { ok: true, bytes: 0 };
}
if (mediaInflightBytes + bytes > MAX_MEDIA_INFLIGHT_BYTES) {
counters.increment("media.inflight_rejected");
const response = errorResponse(
429,
"Too many concurrent media responses in flight; retry.",
"governance_error",
undefined,
"media_inflight",
);
// Set in place, as the governance rate-limit denial does. A rebuild would
// mint a second Response for a header the first one can already carry.
response.headers.set("Retry-After", "1");
return { ok: false, response };
}
mediaInflightBytes += bytes;
return { ok: true, bytes };
}
/** Releases a reservation. Runs in a `finally`, so every exit path - return,
* abort, throw, and the over-cap read - gives the bytes back. */
function releaseMediaInflight(reservation: MediaReservation): void {
if (reservation.ok) {
mediaInflightBytes -= reservation.bytes;
}
}
/**
* Units the provider's parsed body proves it delivered, gated on the PATHNAME
* and never on the body's shape: `/v1/embeddings` answers with `data: []` too,
* so a shape-driven count would bill embeddings as images.
*/
export function countUnits(pathname: string, body: unknown): MediaUnits {
if (pathname === IMAGES_PATH) {
const data = (body as { data?: unknown } | undefined)?.data;
return Array.isArray(data) ? { imageCount: clampCount(data.length) } : {};
}
if (pathname === TRANSCRIPTIONS_PATH) {
const parsed = body as
| { usage?: { seconds?: unknown }; duration?: unknown }
| undefined;
// `usage.seconds` first: a body can carry it without carrying `duration`,
// so taking it first strictly increases the priced population.
const seconds = clampSeconds(parsed?.usage?.seconds) ??
clampSeconds(parsed?.duration);
return seconds === undefined ? {} : { audioSeconds: seconds };
}
return {};
}
/**
* The response body's own token usage, through the one shared normalizer, so a
* media surface that bills on tokens as well as units is counted from the same
* parse as its units.
*
* An all-zero result is reported as ABSENT rather than as zeros: a body that
* carries no token fields must not become a billed zero (decision-log 56), and
* a genuinely zero-token response prices identically either way.
*/
export function tokensFrom(
body: unknown,
): RequestDispatch["tokens"] | undefined {
const carrier = body as
| { usage?: UsageShape; usageMetadata?: UsageShape }
| undefined;
const usage = carrier?.usage ?? carrier?.usageMetadata;
if (!usage || typeof usage !== "object") {
return undefined;
}
const n = normalizeUsage(usage);
if (n.prompt + n.completion + n.cached + n.cacheCreation === 0) {
return undefined;
}
return {
prompt: n.prompt,
completion: n.completion,
cached: n.cached,
cacheCreation: n.cacheCreation,
};
}
/** Parses buffered provider bytes, or undefined when they are not JSON - a
* transcription can legitimately answer `text`, `srt` or `vtt`. The bytes reach
* the client either way, so this must never throw. */
function safeJson(bytes: Uint8Array): unknown {
try {
return JSON.parse(new TextDecoder().decode(bytes));
} catch {
return undefined;
}
}
/**
* Records the provider's own status, from whichever shape it arrived in, and
* reports whether a provider was reached at all.
*
* `rawProxy` is the single `!ok` authority: it THROWS a ProviderError rather
* than returning a `!ok` Response, and every `generateImage` adapter does the
* same. So on a provider error the status arrives as the error, never as
* `upstream.status`, and a route that read only `upstream.status` would leave a
* provider 400 indistinguishable from a refused connection - which is exactly
* the discrimination the write gate depends on.
*
* Residual: an adapter that raises ProviderError from its OWN request validation
* (azure with no deployment, gemini with an unsupported response_format) records
* a status no provider produced. It cannot cause an over-bill, because no
* quantity is ever written on this path.
*/
function recordProviderStatus(
ctx: AppContext,
req: Request,
error: unknown,
): boolean {
if (!(error instanceof ProviderError)) {
return false;
}
if (!mergeRequestStatus(req, error.status)) {
ctx.metrics.increment("accounting.status_dropped");
}
return true;
}
/**
* Dispatches a buffered media surface and writes the accounting channel.
*
* The write gate is the whole point: `units` and `tokens` reach the channel only
* once the provider's response headers are in hand AND `upstream.ok`. A gate at
* the read site could not do this - `mapDispatchError` returns a Response
* carrying the PROVIDER's status, so a provider 400 and a gateway 400 are
* indistinguishable there, and a refused connection produces no Response at
* all. This site knows the provider's status exactly, from `upstream.status` on
* the success path and from the thrown ProviderError otherwise.
*/
async function dispatchBufferedMedia(
ctx: AppContext,
req: Request,
pathname: string,
target: ResolvedTarget,
reservationBytes: number,
dispatch: (signal: AbortSignal) => Promise<Response>,
): Promise<Response> {
const reservation = reserveMediaInflight(reservationBytes, ctx.metrics);
if (!reservation.ok) {
return reservation.response;
}
let headersSeen = false;
try {
setRequestDispatch(
req,
{ providerId: target.providerId, model: target.model },
ctx.metrics,
);
const upstream = await dispatchThenDetach(req, dispatch);
headersSeen = true;
if (!mergeRequestStatus(req, upstream.status)) {
// A dropped status is a silently unbilled row; keep it observable.
ctx.metrics.increment("accounting.status_dropped");
}
if (!upstream.ok) {
return upstream; // no units, no tokens: no provider work was delivered
}
let bytes: Uint8Array;
try {
bytes = await readCappedBytes(upstream.body, reservation.bytes);
} catch (error) {
if (error instanceof GatewayError && error.status === 413) {
// readCappedBytes' 413 describes a REQUEST body. An over-cap PROVIDER
// response is a provider-contract violation, so it must not reach the
// client as a client error - and the row stays unbilled (status
// recorded, no units, no tokens) and counted.
ctx.metrics.increment("media.body_cap_exceeded");
return errorResponse(
502,
`Provider response exceeded the ${reservation.bytes}-byte buffer ` +
`limit for this surface.`,
"provider_error",
undefined,
"media_body_cap_exceeded",
);
}
throw error;
}
const parsed = safeJson(bytes);
mergeRequestUnits(req, countUnits(pathname, parsed));
mergeRequestTokens(req, tokensFrom(parsed));
// readCappedBytes always merges into a fresh `new Uint8Array(total)`, which
// is ArrayBuffer-backed; the assertion only narrows it to what BodyInit
// requires (the same narrowing gemini.ts does for pcmToWav).
return rebuild(upstream, bytes as Uint8Array<ArrayBuffer>);
} catch (error) {
const reached = recordProviderStatus(ctx, req, error);
if (!headersSeen && !reached && req.signal.aborted) {
ctx.metrics.increment("media.abort_pre_headers");
}
return mapDispatchError(error);
} finally {
releaseMediaInflight(reservation);
}
}
export function registerAdvancedRoutes(router: Router, ctx: AppContext): void {
router.post("/v1/embeddings", async (req) => {
const parsed = EmbeddingRequestSchema.safeParse(await parseJsonBody(req));
if (!parsed.success) {
return validationErrorResponse(parsed.error);
}
ctx.metrics.increment("requests.embeddings");
try {
const target = targetFromModel(ctx, parsed.data.model);
requireCapability(target, "supportsEmbeddings", "embeddings");
if (!target.adapter.embeddings) {
throw new GatewayError(
400,
`Provider "${target.providerId}" has no embeddings adapter.`,
);
}
return await target.adapter.embeddings(
{ ...parsed.data, model: target.model },
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
router.post(IMAGES_PATH, async (req) => {
const parsed = ImageGenerationRequestSchema.safeParse(
await parseJsonBody(req),
);
if (!parsed.success) {
return validationErrorResponse(parsed.error);
}
ctx.metrics.increment("requests.images");
// True once the native adapter has been called. Scopes the pre-dispatch
// abort counter below to that branch: the rawProxy branch counts its own
// inside dispatchBufferedMedia, and mapDispatchError rethrows an AbortError
// through here, so an unscoped increment would count it twice.
let nativeDispatched = false;
try {
const target = parsed.data.model
? targetFromModel(ctx, parsed.data.model)
: targetFromQuery(ctx, req);
requireCapability(target, "supportsImages", "image generation");
const body = parsed.data.model
? { ...parsed.data, model: target.model }
: parsed.data;
// Native image adapters (e.g. Google Imagen `:predict`) translate the
// request themselves; OpenAI-wire providers (OpenAI, Nebius) keep the
// byte-identical rawProxy passthrough.
if (target.adapter.generateImage) {
const generate = target.adapter.generateImage.bind(target.adapter);
setRequestDispatch(
req,
{ providerId: target.providerId, model: target.model },
ctx.metrics,
);
nativeDispatched = true;
// dispatchDetached, not dispatchThenDetach: this adapter resolves only
// after it has read and parsed the provider's whole body, so there is
// no header boundary here for the route to detach at. Keeping the
// client's abort live would cancel a render the provider has already
// performed - the images would be charged and none of them counted.
const result = await dispatchDetached(
req,
(signal) => generate(body, { signal }),
);
// Every generateImage adapter throws ProviderError on a non-2xx, so a
// resolved call IS the evidence a provider was reached and delivered.
// The typed return cannot carry the exact code, so 200 stands for it.
// No reservation: the adapter buffers its own JSON, so there is no
// gateway-side capped read for a reservation to bound.
mergeRequestStatus(req, 200);
mergeRequestUnits(req, countUnits(IMAGES_PATH, result));
mergeRequestTokens(req, tokensFrom(result));
return jsonResponse(result);
}
const proxy = requireRawProxy(target);
return await dispatchBufferedMedia(
ctx,
req,
IMAGES_PATH,
target,
mediaReservationBytes(IMAGES_PATH, parsed.data),
(signal) =>
proxy(
"/images/generations",
jsonForward("/images/generations", body),
{ signal },
),
);
} catch (error) {
// The native generateImage branch dispatches inside this try, and its
// adapters throw ProviderError on a non-2xx, so this is where that
// branch's provider status is recovered. dispatchBufferedMedia records
// its own and never lets a ProviderError escape, so there is no overlap.
const reached = recordProviderStatus(ctx, req, error);
if (nativeDispatched && !reached && req.signal.aborted) {
// The only abort this branch can still throw on is one that had already
// landed at entry, so no provider was reached and no quantity exists.
ctx.metrics.increment("media.abort_pre_headers");
}
return mapDispatchError(error);
}
});
router.post(SPEECH_PATH, async (req) => {
const parsed = SpeechRequestSchema.safeParse(await parseJsonBody(req));
if (!parsed.success) {
return validationErrorResponse(parsed.error);
}
ctx.metrics.increment("requests.speech");
// Request-side quantity, so it is free and complete before dispatch - but it
// stays in a LOCAL until the provider's headers prove a provider was
// reached. Code points, not UTF-16 units: `[...s].length` counts an
// astral-plane character once.
const characterCount = clampCount([...parsed.data.input].length);
// Gemini TTS is token-priced and answers with audio bytes, so its usage
// block arrives through this callback rather than through a JSON body. The
// numbers are provider-derived; the route bounds them here.
let reported: RequestDispatch["tokens"] | undefined;
const onUsage = (usage: { prompt?: number; completion?: number }): void => {
reported = {
prompt: clampCount(usage.prompt),
completion: clampCount(usage.completion),
cached: 0,
cacheCreation: 0,
};
};
let headersSeen = false;
try {
const target = targetFromModel(ctx, parsed.data.model);
requireCapability(target, "supportsAudio", "text-to-speech");
const proxy = requireRawProxy(target);
setRequestDispatch(
req,
{ providerId: target.providerId, model: target.model },
ctx.metrics,
);
const upstream = await dispatchThenDetach(req, (signal) =>
proxy(
"/audio/speech",
jsonForward("/audio/speech", { ...parsed.data, model: target.model }),
{ signal, onUsage },
));
headersSeen = true;
if (!mergeRequestStatus(req, upstream.status)) {
ctx.metrics.increment("accounting.status_dropped");
}
if (upstream.ok) {
mergeRequestUnits(req, { characterCount });
mergeRequestTokens(req, reported);
}
// No buffering, no decode, no rebuild, no reservation: the provider's
// audio stream is returned untouched.
return upstream;
} catch (error) {
const reached = recordProviderStatus(ctx, req, error);
if (!headersSeen && !reached && req.signal.aborted) {
ctx.metrics.increment("media.abort_pre_headers");
}
return mapDispatchError(error);
}
});
// Multipart passthrough: routed by ?provider= or the default provider.
router.post(TRANSCRIPTIONS_PATH, async (req) => {
ctx.metrics.increment("requests.transcriptions");
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsAudio", "transcription");
const proxy = requireRawProxy(target);
return await dispatchBufferedMedia(
ctx,
req,
TRANSCRIPTIONS_PATH,
target,
mediaReservationBytes(TRANSCRIPTIONS_PATH),
(signal) => proxy("/audio/transcriptions", req, { signal }),
);
} catch (error) {
return mapDispatchError(error);
}
});
router.post("/v1/files", async (req) => {
ctx.metrics.increment("requests.files");
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "file APIs");
return await requireRawProxy(target)("/files", req, {
signal: req.signal,
});
} catch (error) {
return mapDispatchError(error);
}
});
router.get("/v1/files", async (req) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "file APIs");
return await requireRawProxy(target)("/files", req, {
signal: req.signal,
});
} catch (error) {
return mapDispatchError(error);
}
});
// Retrieve the raw bytes/results of a file (e.g. a batch output_file_id).
// Registered before /v1/files/:id so the 4-segment path wins its own match.
router.get("/v1/files/:id/content", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "file APIs");
return await requireRawProxy(target)(
`/files/${match.pathname.groups.id}/content`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
router.get("/v1/files/:id", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "file APIs");
return await requireRawProxy(target)(
`/files/${match.pathname.groups.id}`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
router.delete("/v1/files/:id", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "file APIs");
return await requireRawProxy(target)(
`/files/${match.pathname.groups.id}`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
router.post("/v1/batches", async (req) => {
ctx.metrics.increment("requests.batches");
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "batch APIs");
return await requireRawProxy(target)("/batches", req, {
signal: req.signal,
});
} catch (error) {
return mapDispatchError(error);
}
});
router.get("/v1/batches/:id", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "batch APIs");
return await requireRawProxy(target)(
`/batches/${match.pathname.groups.id}`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
// Provider-specific batch result stream. Keep this before the generic batch
// retrieval route so `/results` is not mistaken for a batch id.
router.get("/v1/batches/:id/results", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "batch APIs");
return await requireRawProxy(target)(
`/batches/${match.pathname.groups.id}/results`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
router.get("/v1/batches", async (req) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "batch APIs");
return await requireRawProxy(target)("/batches", req, {
signal: req.signal,
});
} catch (error) {
return mapDispatchError(error);
}
});
router.post("/v1/batches/:id/cancel", async (req, match) => {
try {
const target = targetFromQuery(ctx, req);
requireCapability(target, "supportsFiles", "batch APIs");
return await requireRawProxy(target)(
`/batches/${match.pathname.groups.id}/cancel`,
req,
{ signal: req.signal },
);
} catch (error) {
return mapDispatchError(error);
}
});
}
Reference in New Issue View Git Blame Copy Permalink