@mupt-ai/dari-gateway 0.2.0 → 0.3.0

package/dist/providers/provider_call.js

@@ -0,0 +1,296 @@
+import { isRecord, } from "@mupt-ai/dari-router/policy-engine";
+import { executionProviderForModel, keySourceForModel } from "../models/execution_provider.js";
+import { HttpError } from "../http/errors.js";
+import { CustomerDisconnectedError, createProviderAttemptAbort, } from "../runtime/customer_disconnect.js";
+import { isProviderFailureMessage, providerFailureFromCaught, providerFailureFromEvent, } from "../runtime/provider_stream_outcome.js";
+import { providerAuthForSelectedModel } from "../strategies/routing.js";
+import { scrubSecretValues } from "../telemetry/scrub.js";
+import { usageContextFor, } from "../telemetry/usage.js";
+import { awsRoleCredentials } from "./aws_assume_role.js";
+import { attachProviderResponseCapture, createProviderResponseCapture, } from "./response_capture.js";
+import { completeInteractionModel, streamInteractionModel, } from "./pi.js";
+import { promptTokenEstimateFor } from "./cancel_estimate.js";
+export function prepareProviderExecution(args) {
+    const model = args.decision.selectedModel;
+    const promptEstimate = promptTokenEstimateFor(args.costEstimates, model, args.decision.reasoningEffort);
+    const limit = args.credentials.model_configs?.[model]?.standard_context_limit;
+    if (args.credentials.allow_long_context !== true && limit != null && (promptEstimate === undefined || promptEstimate.promptTokens > limit)) {
+        throw new HttpError(400, "The prompt exceeds this model's standard-price context range. Shorten or compact the conversation.", "context_length_exceeded");
+    }
+    const provider = executionProviderForModel(args.credentials, model);
+    const supplied = args.credentials.provider_credentials?.[model];
+    const providerKeySource = keySourceForModel(args.credentials, model);
+    // A strict router serves a plan only when the platform supplied it for this
+    // request, never a pooled or personal one it resolved.
+    if (args.credentials.strict === true && providerKeySource === "subscription" && !supplied) {
+        throw new HttpError(400, `Strict routers do not serve '${model}' from a personal subscription.`, "strict_subscription_refused", "model");
+    }
+    const auth = supplied
+        ? { type: "bearer", apiKey: supplied.apiKey.trim() }
+        : providerAuthForSelectedModel({ selectedModel: model, credentials: args.credentials });
+    return {
+        model,
+        apiKey: auth.type === "bearer" ? auth.apiKey : "",
+        auth,
+        pricing: args.pricing(model),
+        config: args.credentials.model_configs?.[model] ? {
+            ...args.credentials.model_configs[model],
+            allow_long_context: args.credentials.allow_long_context === true,
+            zdr: args.credentials.zdr === true,
+        } : undefined,
+        conversationId: args.conversationId,
+        sessionId: args.promptCacheKey?.trim()
+            ? args.promptCacheKey
+            : args.conversationId,
+        reasoningEffort: args.decision.reasoningEffort,
+        promptEstimate,
+        ...(supplied?.baseUrl || supplied?.headers
+            ? { endpoint: { baseUrl: supplied.baseUrl, headers: supplied.headers } }
+            : {}),
+        provider,
+        providerKeySource,
+        // Only resolved personal subscriptions are held at their usage limit;
+        // the platform manages the plans it supplies.
+        ...(providerKeySource === "subscription" && !supplied && args.credentials.attributed_user_id
+            ? {
+                subscriptionUserId: args.credentials.attributed_user_id,
+                subscriptionCredentialId: args.credentials.subscription_credential_ids?.[provider],
+            }
+            : {}),
+        usageContext: usageContextFor({
+            billable: args.billable,
+            credentials: args.credentials,
+            routerId: args.routerId,
+            selectedModel: model,
+            keySource: providerKeySource,
+            tags: args.tags ?? {},
+        }),
+        headers: {
+            "X-Dari-Selected-Model": model,
+            "X-Dari-Thinking-Level": args.decision.reasoningEffort,
+            "X-Dari-Provider-Key-Source": providerKeySource,
+            "X-Dari-Conversation-Id": args.conversationId,
+            ...(args.decision.leaseTurnsRemaining === undefined
+                ? {}
+                : { "X-Dari-Lease-Turns-Remaining": String(args.decision.leaseTurnsRemaining) }),
+        },
+    };
+}
+export function startProviderCompletion(args) {
+    const customerAbort = createProviderAttemptAbort({
+        clientSignal: args.clientSignal,
+    });
+    const responseCapture = responseCaptureFor(args.execution);
+    const promise = customerAbort.aborted
+        ? Promise.reject(new CustomerDisconnectedError())
+        : completeProvider({
+            interaction: args.interaction,
+            execution: args.execution,
+            signal: customerAbort.signal,
+            fetch: responseCapture?.fetch,
+        });
+    void promise.then(() => customerAbort.dispose(), () => customerAbort.dispose());
+    promise.catch(() => undefined);
+    return { promise, customerAbort, responseCapture };
+}
+function responseCaptureFor(execution) {
+    return execution.providerKeySource === "subscription"
+        ? createProviderResponseCapture()
+        : undefined;
+}
+export async function startProviderStream(args) {
+    const providerStartedMs = Date.now();
+    const customerAbort = createProviderAttemptAbort({
+        clientSignal: args.clientSignal,
+    });
+    if (customerAbort.aborted) {
+        customerAbort.dispose();
+        throw new CustomerDisconnectedError();
+    }
+    const responseCapture = responseCaptureFor(args.execution);
+    try {
+        const providerStream = await streamInteractionModel({
+            model: args.execution.model,
+            request: args.interaction,
+            apiKey: args.execution.apiKey,
+            auth: await awsRoleCredentials.resolve(args.execution.auth),
+            pricing: args.execution.pricing,
+            config: args.execution.config,
+            sessionId: args.execution.sessionId,
+            selectedReasoningEffort: args.execution.reasoningEffort,
+            endpoint: args.execution.endpoint,
+            signal: customerAbort.signal,
+            fetch: responseCapture?.fetch,
+        });
+        return {
+            providerStream: {
+                ...providerStream,
+                events: scrubProviderEvents(providerStream.events, args.execution.auth),
+            },
+            providerStartedMs,
+            customerAbort,
+            responseCapture,
+        };
+    }
+    catch (error) {
+        const canceled = customerAbort.aborted;
+        customerAbort.dispose();
+        if (canceled)
+            throw new CustomerDisconnectedError();
+        throw attachProviderResponseCapture(error, responseCapture);
+    }
+}
+export async function primeProviderStream(events, customerAbort) {
+    const iterator = events[Symbol.asyncIterator]();
+    const prefix = [];
+    try {
+        let result = await iterator.next();
+        while (!result.done) {
+            if (customerAbort.aborted)
+                throw new CustomerDisconnectedError();
+            const event = result.value;
+            const failure = initialStreamFailure(event);
+            if (failure) {
+                await closeAsyncIterator(iterator);
+                return { outcome: "error", failure };
+            }
+            prefix.push(event);
+            if (isMeaningfulStreamEvent(event)) {
+                return continueWithPrefix(prefix, iterator);
+            }
+            result = await iterator.next();
+        }
+    }
+    catch (error) {
+        if (customerAbort.aborted) {
+            void closeAsyncIterator(iterator);
+            throw new CustomerDisconnectedError();
+        }
+        await closeAsyncIterator(iterator);
+        throw error;
+    }
+    if (customerAbort.aborted) {
+        void closeAsyncIterator(iterator);
+        throw new CustomerDisconnectedError();
+    }
+    await closeAsyncIterator(iterator);
+    return {
+        outcome: "error",
+        failure: providerFailureFromCaught("Provider stream ended without a terminal event."),
+    };
+}
+export function scrubProviderMessage(message, auth) {
+    const scrubbed = scrubSecretValues(message, providerSecrets(auth));
+    if (!isRecord(scrubbed)) {
+        throw new Error("Scrubbing a provider message changed its object shape.");
+    }
+    return scrubbed;
+}
+async function completeProvider(args) {
+    const providerStartedAt = new Date();
+    const completion = await completeInteractionModel({
+        model: args.execution.model,
+        request: args.interaction,
+        apiKey: args.execution.apiKey,
+        auth: await awsRoleCredentials.resolve(args.execution.auth),
+        pricing: args.execution.pricing,
+        config: args.execution.config,
+        sessionId: args.execution.sessionId,
+        selectedReasoningEffort: args.execution.reasoningEffort,
+        endpoint: args.execution.endpoint,
+        signal: args.signal,
+        fetch: args.fetch,
+    });
+    return {
+        ...completion,
+        message: scrubProviderMessage(completion.message, args.execution.auth),
+        providerLatencyMs: Date.now() - providerStartedAt.getTime(),
+        providerStartedAt,
+    };
+}
+function initialStreamFailure(event) {
+    if (event.type === "error")
+        return providerFailureFromEvent(event);
+    if (event.type === "done" &&
+        isRecord(event.message) &&
+        isProviderFailureMessage(event.message)) {
+        return event.message;
+    }
+    return null;
+}
+function isMeaningfulStreamEvent(event) {
+    if (event.type === "done")
+        return true;
+    if (event.type === "text_delta" || event.type === "thinking_delta") {
+        return typeof event.delta === "string" && event.delta.length > 0;
+    }
+    if (event.type === "toolcall_start" ||
+        event.type === "toolcall_delta" ||
+        event.type === "toolcall_end" ||
+        event.type === "server_tool_call_start" ||
+        event.type === "server_tool_call_end" ||
+        event.type === "tool_search_result_start" ||
+        event.type === "tool_search_result_end") {
+        return true;
+    }
+    if (event.type !== "provider_event" || !isRecord(event.event))
+        return false;
+    const type = event.event.type;
+    if (typeof type !== "string")
+        return false;
+    if (type.startsWith("response.web_search_call."))
+        return true;
+    return type === "response.output_item.added" &&
+        isRecord(event.event.item) && event.event.item.type === "web_search_call";
+}
+function continueWithPrefix(prefix, iterator) {
+    return {
+        outcome: "continue",
+        events: {
+            async *[Symbol.asyncIterator]() {
+                let exhausted = false;
+                try {
+                    yield* prefix;
+                    let result = await iterator.next();
+                    while (!result.done) {
+                        yield result.value;
+                        result = await iterator.next();
+                    }
+                    exhausted = true;
+                }
+                finally {
+                    if (!exhausted)
+                        await closeAsyncIterator(iterator);
+                }
+            },
+        },
+    };
+}
+async function closeAsyncIterator(iterator) {
+    try {
+        await iterator.return?.();
+    }
+    catch {
+        // Cleanup must not replace the provider failure or client cancellation.
+    }
+}
+function providerSecrets(auth) {
+    if (auth.type === "bearer")
+        return [auth.apiKey];
+    if (auth.type === "aws_assume_role")
+        return awsRoleCredentials.secrets(auth);
+    return [
+        auth.access_key_id,
+        auth.secret_access_key,
+        ...(auth.session_token ? [auth.session_token] : []),
+    ];
+}
+function scrubProviderEvents(events, auth) {
+    return {
+        async *[Symbol.asyncIterator]() {
+            for await (const event of events) {
+                yield scrubProviderMessage(event, auth);
+            }
+        },
+    };
+}

package/dist/providers/provider_reported_cost.d.ts

@@ -0,0 +1,10 @@
+import type { ProviderFetch } from "./response_capture.js";
+export declare function reportsProviderCost(model: Record<string, unknown>): boolean;
+export type ProviderReportedCostTap = {
+    fetch: ProviderFetch;
+    cost: () => number | undefined;
+};
+export declare function providerReportedCostTap(fetch?: ProviderFetch): ProviderReportedCostTap;
+export declare function messageWithProviderReportedCost(message: Record<string, unknown>, cost: number | undefined): Record<string, unknown>;
+export declare function eventsWithProviderReportedCost(events: AsyncIterable<Record<string, unknown>>, tap: ProviderReportedCostTap): AsyncIterable<Record<string, unknown>>;
+export declare function reportedCost(body: unknown): number | undefined;

package/dist/providers/provider_reported_cost.js

@@ -0,0 +1,123 @@
+import { isRecord } from "@mupt-ai/dari-router/policy-engine";
+import { gatewayFor } from "../models/gateways.js";
+import { providerFailureFromCaught } from "../runtime/provider_stream_outcome.js";
+// Gateways bill per call and say what they billed: OpenRouter as `usage.cost`
+// on the final chunk, Vercel AI Gateway as `providerMetadata.gateway.cost`.
+// Pi prices usage from model rates and drops both, so the amount is read off
+// the raw response bytes and replaces the priced total.
+export function reportsProviderCost(model) {
+    return gatewayFor(model.provider) !== undefined;
+}
+export function providerReportedCostTap(fetch = globalThis.fetch) {
+    let cost;
+    return {
+        cost: () => cost,
+        fetch: async (input, init) => {
+            const response = await fetch(input, init);
+            if (!response.ok || response.body === null)
+                return response;
+            cost = undefined;
+            return new Response(response.body.pipeThrough(costScanner((value) => {
+                cost = value;
+            })), { status: response.status, statusText: response.statusText, headers: response.headers });
+        },
+    };
+}
+export function messageWithProviderReportedCost(message, cost) {
+    if (cost === undefined)
+        return message;
+    // A thrown error (or other record without a stop reason) becomes a failure
+    // message first: spreading an Error drops its non-enumerable `message`, and
+    // failure handling keeps usage only on messages that stopped in error.
+    const base = message.stopReason === undefined
+        ? { ...providerFailureFromCaught(message), ...(isRecord(message.usage) ? { usage: message.usage } : {}) }
+        : message;
+    const usage = isRecord(base.usage) ? base.usage : {};
+    return {
+        ...base,
+        usage: {
+            ...usage,
+            cost: { ...(isRecord(usage.cost) ? usage.cost : {}), total: cost },
+            costSource: "provider",
+        },
+    };
+}
+// Aborted and failed calls end in an error event carrying the partial
+// message; it keeps whatever cost the provider reported before the end.
+export async function* eventsWithProviderReportedCost(events, tap) {
+    for await (const event of events) {
+        if (event.type === "done" && isRecord(event.message)) {
+            yield { ...event, message: messageWithProviderReportedCost(event.message, tap.cost()) };
+            continue;
+        }
+        if (event.type === "error" && isRecord(event.error)) {
+            yield { ...event, error: messageWithProviderReportedCost(event.error, tap.cost()) };
+            continue;
+        }
+        yield event;
+    }
+}
+export function reportedCost(body) {
+    if (!isRecord(body))
+        return undefined;
+    const gateway = isRecord(body.providerMetadata) && isRecord(body.providerMetadata.gateway)
+        ? body.providerMetadata.gateway
+        : undefined;
+    return usd(isRecord(body.usage) ? body.usage.cost : undefined) ?? usd(gateway?.cost);
+}
+function usd(value) {
+    const amount = typeof value === "string" && value.trim() !== "" ? Number(value) : value;
+    return typeof amount === "number" && Number.isFinite(amount) && amount >= 0 ? amount : undefined;
+}
+// Reads SSE `data:` lines and plain JSON bodies alike; the last reported
+// amount wins.
+function costScanner(onCost) {
+    const decoder = new TextDecoder();
+    let pending = "";
+    let body = "";
+    let sse = false;
+    const scan = (line) => {
+        if (!line.startsWith("data:") || !line.includes("cost"))
+            return;
+        try {
+            const cost = reportedCost(JSON.parse(line.slice("data:".length)));
+            if (cost !== undefined)
+                onCost(cost);
+        }
+        catch {
+            // Not JSON: [DONE] or a keep-alive.
+        }
+    };
+    return new TransformStream({
+        transform(chunk, controller) {
+            controller.enqueue(chunk);
+            const text = decoder.decode(chunk, { stream: true });
+            pending += text;
+            const lines = pending.split("\n");
+            pending = lines.pop() ?? "";
+            for (const line of lines) {
+                const trimmed = line.trimEnd();
+                if (trimmed.startsWith("data:"))
+                    sse = true;
+                scan(trimmed);
+            }
+            if (!sse && body.length < 1024 * 1024)
+                body += text;
+        },
+        flush() {
+            const tail = decoder.decode();
+            scan((pending + tail).trimEnd());
+            body += tail;
+            if (sse || !body.includes("cost"))
+                return;
+            try {
+                const cost = reportedCost(JSON.parse(body));
+                if (cost !== undefined)
+                    onCost(cost);
+            }
+            catch {
+                // Not a JSON body.
+            }
+        },
+    });
+}

package/dist/providers/reasoning.d.ts

@@ -0,0 +1,12 @@
+import type { ReasoningEffort } from "@mupt-ai/dari-router/policy-engine";
+import type { InteractionRequest } from "../interaction/types.js";
+type ReasoningRequest = Pick<InteractionRequest, "reasoning">;
+export declare function effectiveReasoningForModel(args: {
+    request: ReasoningRequest;
+    selectedModel: string;
+    supportedLevels: string[];
+    selectedReasoningEffort?: ReasoningEffort;
+    defaultReasoningEffort?: ReasoningEffort;
+}): ReasoningEffort | undefined;
+export declare function validateReasoningEffortForModel(effort: ReasoningEffort, selectedModel: string, supportedLevels: string[]): void;
+export {};

package/dist/providers/reasoning.js

@@ -0,0 +1,30 @@
+import { HttpError } from "../http/errors.js";
+function requestedReasoningEffort(request) {
+    return request.reasoning?.effort;
+}
+export function effectiveReasoningForModel(args) {
+    const requested = requestedReasoningEffort(args.request) ?? args.selectedReasoningEffort;
+    const effective = requested ??
+        args.defaultReasoningEffort ??
+        implicitReasoningEffort(args.supportedLevels);
+    if (effective === undefined)
+        return undefined;
+    validateReasoningEffortForModel(effective, args.selectedModel, args.supportedLevels);
+    return effective;
+}
+export function validateReasoningEffortForModel(effort, selectedModel, supportedLevels) {
+    if (supportedLevels.includes(effort))
+        return;
+    throw new HttpError(400, `reasoning_effort '${effort}' is not supported by ${selectedModel}. Supported values: ${supportedLevels.join(", ")}.`, "unsupported_reasoning_effort", "reasoning_effort");
+}
+// Routed serving always supplies the selected pair's effort. For direct
+// internal provider calls that do not (currently the selector), use cataloged
+// provider defaults when available, then fall back to medium or the closest
+// supported level for custom models. Explicit "off" remains available for
+// routed candidates that support a true non-reasoning mode.
+const IMPLICIT_REASONING_PREFERENCE = ["medium", "low", "high", "minimal", "xhigh", "max"];
+function implicitReasoningEffort(supportedLevels) {
+    // Models with no non-off level (non-reasoning models) keep omitting the
+    // reasoning field entirely.
+    return IMPLICIT_REASONING_PREFERENCE.find((level) => supportedLevels.includes(level));
+}

package/dist/providers/response_capture.d.ts

@@ -0,0 +1,23 @@
+export type CapturedProviderResponse = {
+    status: number;
+    headers: Headers;
+    bodyText: string;
+};
+export type ProviderFetch = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;
+export type ProviderResponseCapture = {
+    fetch: ProviderFetch;
+    last: CapturedProviderResponse | undefined;
+};
+/**
+ * A fetch wrapper that remembers the last failed provider response. pi-ai
+ * flattens provider errors into one message string before they reach the
+ * runtime, so this is the only place the status, rate-limit headers, and
+ * structured error body of a rejected call can still be read.
+ */
+export declare function createProviderResponseCapture(): ProviderResponseCapture;
+/**
+ * A provider call that throws before returning a stream loses its capture with
+ * the started state, so the error carries it to the attempt runner instead.
+ */
+export declare function attachProviderResponseCapture(error: unknown, capture: ProviderResponseCapture | undefined): unknown;
+export declare function providerResponseCaptureFromError(error: unknown): ProviderResponseCapture | undefined;

package/dist/providers/response_capture.js

@@ -0,0 +1,47 @@
+const CAPTURED_BODY_MAX_BYTES = 64 * 1024;
+/**
+ * A fetch wrapper that remembers the last failed provider response. pi-ai
+ * flattens provider errors into one message string before they reach the
+ * runtime, so this is the only place the status, rate-limit headers, and
+ * structured error body of a rejected call can still be read.
+ */
+export function createProviderResponseCapture() {
+    const capture = {
+        last: undefined,
+        fetch: async (input, init) => {
+            const response = await globalThis.fetch(input, init);
+            if (response.ok)
+                return response;
+            const bodyText = await readBodyPrefix(response.clone());
+            capture.last = { status: response.status, headers: response.headers, bodyText };
+            return response;
+        },
+    };
+    return capture;
+}
+async function readBodyPrefix(response) {
+    try {
+        const text = await response.text();
+        return text.length > CAPTURED_BODY_MAX_BYTES ? text.slice(0, CAPTURED_BODY_MAX_BYTES) : text;
+    }
+    catch {
+        return "";
+    }
+}
+const CAPTURE_KEY = Symbol("providerResponseCapture");
+/**
+ * A provider call that throws before returning a stream loses its capture with
+ * the started state, so the error carries it to the attempt runner instead.
+ */
+export function attachProviderResponseCapture(error, capture) {
+    if (capture !== undefined && typeof error === "object" && error !== null) {
+        error[CAPTURE_KEY] = capture;
+    }
+    return error;
+}
+export function providerResponseCaptureFromError(error) {
+    if (typeof error !== "object" || error === null)
+        return undefined;
+    const capture = error[CAPTURE_KEY];
+    return capture === undefined ? undefined : capture;
+}

package/dist/providers/stream_finalization.d.ts

@@ -0,0 +1,46 @@
+import type { ModelPricing, RoutingDecision } from "@mupt-ai/dari-router/policy-engine";
+import type { InteractionRequest, ProviderIdentity } from "../interaction/types.js";
+import type { ProviderStreamOutcome } from "../runtime/provider_stream_outcome.js";
+import type { RoutedProtocol } from "../runtime/routed_protocol.js";
+import { type RoutedRuntimeReporter } from "../runtime/runtime_reporter.js";
+import type { ProviderKeySource } from "../types.js";
+import { type RoutingUsageContext } from "../telemetry/usage.js";
+import { type DeliveredOutput, type PromptTokenEstimate } from "./cancel_estimate.js";
+type StreamExecutionFacts = {
+    provider: string;
+    model: string;
+    providerKeySource: ProviderKeySource;
+    conversationId: string;
+    usageContext: RoutingUsageContext | undefined;
+    pricing: ModelPricing | null;
+    promptEstimate: PromptTokenEstimate | undefined;
+};
+type CreateStreamFinalizerArgs = {
+    requestId: string;
+    interaction: InteractionRequest;
+    decision: RoutingDecision;
+    execution: StreamExecutionFacts;
+    reporter: RoutedRuntimeReporter;
+    protocol: RoutedProtocol;
+    providerLatencyMs: () => number;
+    role: "primary" | "fallback";
+    attempt: number;
+    reportBarrier?: Promise<unknown>;
+    delivered: () => DeliveredOutput;
+    source: ProviderIdentity;
+    startedAtMs: number;
+    cancelBarrierCapMs?: number;
+};
+/**
+ * A client cancel runs its finalization in the background, where a slow
+ * selector would only delay the record and, at shutdown, eat the drain budget.
+ * The record still gets the selector's run when it arrives within the cap;
+ * otherwise it is written without it. The response-owned path is not capped.
+ */
+export declare const CANCEL_REPORT_BARRIER_CAP_MS = 2000;
+export type StreamFinalizer = {
+    beforeResponseClose(outcome: ProviderStreamOutcome): Promise<void>;
+    afterResponseDetached(outcome: ProviderStreamOutcome): Promise<void>;
+};
+export declare function createStreamFinalizer(args: CreateStreamFinalizerArgs): StreamFinalizer;
+export {};