@mupt-ai/dari-gateway 0.2.0 → 0.2.1

package/dist/runtime/routed_runtime.js

@@ -0,0 +1,535 @@
+import { randomUUID } from "node:crypto";
+import { isRecord, sameRoutingCandidate, } from "@mupt-ai/dari-router/policy-engine";
+import { buildRoutingContext } from "../context/routing_context.js";
+import { completeNonStreaming, executeNonStreaming, executeStreaming, prepareStreamingExecution, } from "../providers/execution.js";
+import { prepareProviderExecution, startProviderCompletion, startProviderStream, } from "../providers/provider_call.js";
+import { routeDecision } from "../strategies/route_decision.js";
+import { executionProviderForModel } from "../models/execution_provider.js";
+import { trackBackgroundTelemetry } from "../telemetry/background.js";
+import { estimateBenchmarkReplay } from "../telemetry/benchmark_replay.js";
+import { buildCompletedPrefixHit } from "../telemetry/prefix_state.js";
+import { scrubSecrets } from "../telemetry/scrub.js";
+import { auxiliaryDecision } from "./claude_code/auxiliary_routing.js";
+import { explicitModelDecision, strictModelDecision } from "./explicit_model.js";
+import { clearPendingLease, committedLeaseDecision, leaseHoldDecision, leasePrefetchDue, leasedDecision, pendingLeaseDecision, prefetchNextLease, selectorNegotiatesLeases, } from "./routing_lease.js";
+import { createExperimentRoutedRuntimeReporter, } from "./runtime_reporter.js";
+import { HttpError } from "../http/errors.js";
+export class RoutedRuntimeProviderSettledError extends Error {
+    name = "RoutedRuntimeProviderSettledError";
+    runtimeResult;
+    cause;
+    constructor(message, runtimeResult, cause) {
+        super(message);
+        this.runtimeResult = runtimeResult;
+        this.cause = cause;
+    }
+}
+export { createExperimentRoutedRuntimeReporter } from "./runtime_reporter.js";
+export async function runPreparedRoutedCompletion(args) {
+    const plan = await prepareRuntimeExecution({
+        prepared: args.prepared,
+        selectorBackend: args.selectorBackend,
+        clientSignal: args.clientSignal,
+    });
+    const { prepared } = args;
+    if (!prepared.interaction.stream) {
+        return executeNonStreaming({
+            requestId: prepared.requestId,
+            interaction: prepared.interaction,
+            primary: plan.primary,
+            primaryRetries: plan.primaryRetries,
+            resolveFallback: plan.resolveFallback,
+            reporter: prepared.reporter,
+            protocol: prepared.protocol,
+            startedCompletion: plan.startedCompletion,
+            reportBarrier: plan.reportBarrier,
+            clientSignal: args.clientSignal,
+        });
+    }
+    const streaming = await prepareStreamingExecution({
+        requestId: prepared.requestId,
+        interaction: prepared.interaction,
+        primary: plan.primary,
+        primaryRetries: plan.primaryRetries,
+        resolveFallback: plan.resolveFallback,
+        reporter: prepared.reporter,
+        startedStream: plan.startedStream,
+        initialError: plan.initialStreamError,
+        reportBarrier: plan.reportBarrier,
+        clientSignal: args.clientSignal,
+    });
+    return executeStreaming({
+        requestId: prepared.requestId,
+        interaction: prepared.interaction,
+        choice: streaming.choice,
+        reporter: prepared.reporter,
+        protocol: prepared.protocol,
+        startedStream: streaming.started,
+        attempt: streaming.attempt,
+        failedSetupLatencyMs: streaming.failedSetupLatencyMs,
+        reportBarrier: plan.reportBarrier,
+        clientSignal: args.clientSignal,
+    });
+}
+export async function runPreparedRoutedInteraction(args) {
+    if (args.prepared.interaction.stream) {
+        throw new HttpError(400, "In-process routed runtime results only support non-streaming interactions.", "streaming_not_supported");
+    }
+    const plan = await prepareRuntimeExecution({
+        ...args,
+        requireSelectorDecision: true,
+    });
+    const result = await completeNonStreaming({
+        requestId: args.prepared.requestId,
+        interaction: args.prepared.interaction,
+        primary: plan.primary,
+        primaryRetries: plan.primaryRetries,
+        resolveFallback: plan.resolveFallback,
+        reporter: args.prepared.reporter,
+        protocol: args.prepared.protocol,
+        startedCompletion: plan.startedCompletion,
+        reportBarrier: plan.reportBarrier,
+    });
+    const routedResult = runtimeResult(result, args.prepared);
+    if (result.reportBarrierError !== undefined) {
+        throw new RoutedRuntimeProviderSettledError(errorMessage(result.reportBarrierError), routedResult, result.reportBarrierError);
+    }
+    return routedResult;
+}
+export async function runExplicitRoutedInteraction(args) {
+    const prepared = await prepareExplicitRoutedCompletion(args);
+    return runPreparedRoutedInteraction({
+        prepared,
+        selectorBackend: args.selectorBackend,
+    });
+}
+async function prepareExplicitRoutedCompletion(args) {
+    const receivedAt = new Date();
+    const requestId = randomUUID();
+    const routing = await buildRoutingContext({
+        routerId: args.routerId,
+        interaction: args.interaction,
+        credentials: args.credentials,
+        previousDecision: args.previousDecision,
+        receivedAt,
+    });
+    return {
+        requestId,
+        routerId: args.routerId,
+        receivedAt,
+        interaction: args.interaction,
+        protocol: inProcessProtocol,
+        credentials: args.credentials,
+        billable: args.billable ?? false,
+        tags: {},
+        harnessConversationIdentity: undefined,
+        routing,
+        reporter: args.reporter ?? createExperimentRoutedRuntimeReporter(),
+    };
+}
+async function prepareRuntimeExecution(args) {
+    const { prepared } = args;
+    // A strict router never selects, leases, or falls back: the pinned pair or
+    // an error.
+    if (prepared.credentials.strict === true) {
+        const pinned = strictModelDecision(prepared);
+        prepared.reporter.attachServingDecision(pinned.decision);
+        prepared.reporter.attachDecisionDetails({
+            decision_source: "strict_model",
+            explicit_model: {
+                model: pinned.decision.selectedModel,
+                reasoning_effort: pinned.decision.reasoningEffort,
+            },
+        });
+        return selectedExecutionPlan(prepared, pinned);
+    }
+    // A requested `dari/model/<id>` is served as chosen, before any lease.
+    const explicit = args.requireSelectorDecision ? null : explicitModelDecision(prepared);
+    if (explicit) {
+        prepared.reporter.attachServingDecision(explicit.decision);
+        prepared.reporter.attachDecisionDetails({
+            decision_source: "explicit_model",
+            explicit_model: {
+                model: explicit.decision.selectedModel,
+                reasoning_effort: explicit.decision.reasoningEffort,
+            },
+        });
+        return selectedExecutionPlan(prepared, explicit);
+    }
+    // Live background work bypasses task selection; explicit experiments
+    // still require a selector observation and never consume live leases.
+    const auxiliary = args.requireSelectorDecision ? null : auxiliaryDecision(prepared);
+    if (auxiliary) {
+        prepared.reporter.attachServingDecision(auxiliary.decision);
+        prepared.reporter.attachDecisionDetails({
+            decision_source: "auxiliary_request",
+            auxiliary_request: {
+                request_class: prepared.interaction.agentContext?.requestClass,
+                model: auxiliary.decision.selectedModel,
+                reasoning_effort: auxiliary.decision.reasoningEffort,
+            },
+        });
+        return selectedExecutionPlan(prepared, auxiliary);
+    }
+    // Persisted policy leases (negotiated {action, turns}) take precedence.
+    // Explicit in-process runs are selector experiments: they must observe the
+    // selector and must not consume or prefetch live conversations' leases.
+    const held = args.requireSelectorDecision ? null : leaseHoldDecision(prepared);
+    if (held) {
+        // Serving-only attach: nextDecision stays unset so if a fallback serves
+        // this turn, execution's attachServingDecision moves both the prefix's
+        // next_model and inherited countdown onto the healthy serving pair.
+        prepared.reporter.attachServingDecision(held);
+        prepared.reporter.attachDecisionDetails({
+            decision_source: "lease_hold",
+            lease_hold: {
+                model: held.selectedModel,
+                reasoning_effort: held.reasoningEffort,
+                turns_remaining_after: held.leaseTurnsRemaining ?? 0,
+            },
+        });
+        // Prefetch fires on the turn that leaves lag turns of lease including
+        // itself — training stashes at turnIndex === endTurn - lag, and this
+        // turn's pre-decrement countdown equals the lag exactly then. Firing
+        // unconditionally (rather than skipping when a pending lease exists)
+        // means a stale pending from a failed consume is overwritten with a
+        // fresh decision before the next expiry can re-adopt it.
+        if (leasePrefetchDue(prepared)) {
+            trackBackgroundTelemetry(prefetchNextLease(prepared, args.selectorBackend));
+        }
+        return selectedExecutionPlan(prepared, committedLeaseDecision(prepared, held));
+    }
+    const adopted = args.requireSelectorDecision ? null : pendingLeaseDecision(prepared);
+    if (adopted) {
+        prepared.reporter.attachServingDecision(adopted);
+        prepared.reporter.attachDecisionDetails({
+            decision_source: "prefetched_lease",
+            selector_lease_turns: prepared.routing.pendingLease?.turns ?? null,
+            lease_prefetch_adopted: {
+                model: adopted.selectedModel,
+                reasoning_effort: adopted.reasoningEffort,
+                turns_remaining_after: adopted.leaseTurnsRemaining ?? 0,
+            },
+        });
+        prepared.reporter.attachAdoptedSelectorOutput({
+            outputText: prepared.routing.pendingLease?.outputText ?? null,
+            reason: adopted.reason,
+            candidates: prepared.routing.candidateResolution.candidates,
+        });
+        // Consume the prefetched decision so a later expiry cannot re-adopt it.
+        // Awaited with a cap: the clear is one small internal write, and losing
+        // it silently would leave the stale lease adoptable. Even if the capped
+        // wait elapses, the next lease's own prefetch overwrites the stale rows
+        // before its expiry (the menu's shortest lease exceeds the lag).
+        await clearPendingLease(prepared);
+        return selectedExecutionPlan(prepared, committedLeaseDecision(prepared, adopted));
+    }
+    // The stateless ten-turn hold covers per-turn policies. A lease-negotiating
+    // backend must reach the selector at expiry — the modulo hold would mask it
+    // — so the negotiated protocol bypasses it entirely.
+    const leased = args.requireSelectorDecision || (await selectorNegotiatesLeases(args.selectorBackend))
+        ? null
+        : leasedDecision(prepared);
+    const speculative = leased
+        ? null
+        : await startSpeculativeExecution(prepared, args.selectorBackend, args.clientSignal);
+    if (!speculative) {
+        const selected = leased
+            ?? await selectDecision(prepared, null, args.selectorBackend);
+        if (!leased) {
+            prepared.reporter.attachDecisionDetails({
+                decision_source: "selector",
+                selector_lease_turns: selected.decision.leaseTurnsRemaining === undefined
+                    ? null
+                    : selected.decision.leaseTurnsRemaining + 1,
+            });
+        }
+        return selectedExecutionPlan(prepared, selected);
+    }
+    prepared.reporter.attachDecision(speculative.decision);
+    const primary = {
+        decision: speculative.decision,
+        execution: speculative.execution,
+        role: "primary",
+    };
+    const selectionPromise = selectDecision(prepared, speculative.decision, args.selectorBackend, args.requireSelectorDecision === true);
+    // The execution layer awaits selection as its report barrier only after the
+    // provider settles, which can be long after a rejection. Mark the promise
+    // handled now; fallback still consumes the original and its rejection.
+    selectionPromise.catch(() => undefined);
+    let fallbackPromise = null;
+    const resolveFallback = () => {
+        fallbackPromise ??= selectionPromise.then((selected) => {
+            const decision = fallbackForPrimary(prepared, primary.decision, selected);
+            if (decision)
+                prepared.reporter.attachFallbackDecision(decision);
+            return decision ? providerChoice(prepared, decision, "fallback") : null;
+        });
+        return fallbackPromise;
+    };
+    return {
+        primary,
+        primaryRetries: configuredPrimaryRetries(prepared.credentials.primary_retries),
+        resolveFallback,
+        reportBarrier: selectionPromise,
+        startedCompletion: speculative.transport === "completion" ? speculative.started : undefined,
+        startedStream: speculative.transport === "stream" ? speculative.started : undefined,
+        initialStreamError: speculative.transport === "stream"
+            ? speculative.initialError
+            : undefined,
+    };
+}
+function selectedExecutionPlan(prepared, selected) {
+    return {
+        primary: providerChoice(prepared, selected.decision, "primary"),
+        primaryRetries: configuredPrimaryRetries(prepared.credentials.primary_retries),
+        resolveFallback: async () => {
+            if (!selected.fallbackDecision)
+                return null;
+            const fallbackDecision = selected.decision.leaseTurnsRemaining === undefined
+                ? selected.fallbackDecision
+                : {
+                    ...selected.fallbackDecision,
+                    leaseTurnsRemaining: selected.decision.leaseTurnsRemaining,
+                };
+            return providerChoice(prepared, fallbackDecision, "fallback");
+        },
+    };
+}
+async function startSpeculativeExecution(prepared, selectorBackend, clientSignal) {
+    // Speculation is a new dispatch. A client already gone during request
+    // preparation must never cause fresh paid provider work.
+    if (clientSignal?.aborted === true)
+        return null;
+    const { credentials, interaction, routing } = prepared;
+    const resolution = routing.candidateResolution;
+    if (credentials.speculative_routing !== true ||
+        selectorBackend.supportsSpeculativeRouting === false ||
+        credentials.credit_exhausted ||
+        !routing.previousDecision) {
+        return null;
+    }
+    const previous = {
+        model: routing.previousDecision.model,
+        reasoningEffort: routing.previousDecision.reasoningEffort,
+    };
+    if (!resolution.candidates.some((candidate) => sameRoutingCandidate(candidate, previous))) {
+        return null;
+    }
+    const decision = {
+        selectedModel: previous.model,
+        reasoningEffort: previous.reasoningEffort,
+        reason: routing.previousDecision.reason ||
+            "Continuing the previous model while selecting the next route.",
+    };
+    try {
+        const execution = providerExecution(prepared, decision);
+        if (interaction.stream) {
+            try {
+                const started = await startProviderStream({
+                    interaction,
+                    execution,
+                    clientSignal,
+                });
+                return { transport: "stream", decision, execution, started };
+            }
+            catch (initialError) {
+                return { transport: "stream", decision, execution, initialError };
+            }
+        }
+        const started = startProviderCompletion({
+            interaction,
+            execution,
+            clientSignal,
+        });
+        return { transport: "completion", decision, execution, started };
+    }
+    catch (error) {
+        console.error("speculative provider setup failed", scrubSecrets(error instanceof Error ? error.message : String(error)));
+        return null;
+    }
+}
+async function selectDecision(prepared, servingDecision, selectorBackend, requireSelectorDecision = false) {
+    const { credentials, reporter, routing } = prepared;
+    try {
+        const routed = await routeDecision({
+            selectorMessages: routing.selectorMessages,
+            selectorImages: routing.selectorImages,
+            resolution: routing.candidateResolution,
+            credentials,
+            previousDecision: routing.previousDecision,
+            conversationId: routing.conversationId,
+            selectorBackend,
+            leaseHistory: routing.leaseHistory,
+        });
+        if (routed.details)
+            reporter.attachDecisionDetails(routed.details);
+        if (servingDecision) {
+            if (!routed.run) {
+                throw new Error("Speculative routing completed without a selector run.");
+            }
+            reporter.attachSpeculativeSelectorRun(routed.run, servingDecision);
+            reporter.attachDecisionDetails({
+                decision_source: "speculative_selector",
+                selector_lease_turns: routed.decision.leaseTurnsRemaining === undefined
+                    ? null
+                    : routed.decision.leaseTurnsRemaining + 1,
+                speculative_routing: {
+                    serving_model: servingDecision.selectedModel,
+                    serving_reasoning_effort: servingDecision.reasoningEffort,
+                    next_model: routed.decision.selectedModel,
+                    next_reasoning_effort: routed.decision.reasoningEffort,
+                },
+            });
+        }
+        else if (routed.run) {
+            reporter.attachSelectorRun(routed.run);
+        }
+        else {
+            reporter.attachDecision(routed.decision);
+            reporter.attachFallbackDecision(routed.fallbackDecision);
+        }
+        return routed;
+    }
+    catch (error) {
+        const details = isRecord(error) && isRecord(error.decisionDetails)
+            ? error.decisionDetails
+            : null;
+        if (details)
+            reporter.attachDecisionDetails(details);
+        if (!servingDecision) {
+            reporter.reportSelectorFailure(error);
+            throw error;
+        }
+        if (requireSelectorDecision) {
+            throw error;
+        }
+        const code = error instanceof HttpError ? error.code : "selector_execution_failed";
+        console.error("speculative routing selector failed", JSON.stringify({
+            code,
+            message: scrubSecrets(error instanceof Error ? error.message : String(error)),
+        }));
+        reporter.attachDecision(servingDecision);
+        reporter.attachDecisionDetails({
+            decision_source: "speculative_selector",
+            speculative_routing: {
+                serving_model: servingDecision.selectedModel,
+                serving_reasoning_effort: servingDecision.reasoningEffort,
+                next_model: servingDecision.selectedModel,
+                next_reasoning_effort: servingDecision.reasoningEffort,
+                selector_error_code: code,
+            },
+        });
+        return { decision: servingDecision };
+    }
+}
+function errorMessage(error) {
+    return error instanceof Error ? error.message : String(error);
+}
+function configuredPrimaryRetries(value) {
+    const retries = value ?? 1;
+    return Number.isInteger(retries) && retries >= 0 && retries <= 5 ? retries : 1;
+}
+function fallbackForPrimary(prepared, primary, selected) {
+    if (prepared.credentials.model_fallback_enabled === false)
+        return undefined;
+    const fallback = [selected.fallbackDecision, selected.decision].find((candidate) => candidate !== undefined &&
+        candidate.selectedModel !== primary.selectedModel);
+    if (!fallback || selected.decision.leaseTurnsRemaining === undefined) {
+        return fallback;
+    }
+    return {
+        ...fallback,
+        leaseTurnsRemaining: selected.decision.leaseTurnsRemaining,
+    };
+}
+function providerChoice(prepared, decision, role) {
+    return {
+        decision,
+        execution: providerExecution(prepared, decision),
+        role,
+    };
+}
+function providerExecution(prepared, decision) {
+    return prepareProviderExecution({
+        decision,
+        credentials: prepared.credentials,
+        billable: prepared.billable,
+        routerId: prepared.routerId,
+        pricing: prepared.routing.pricing,
+        conversationId: prepared.routing.conversationId,
+        costEstimates: prepared.routing.reportingCostEstimates,
+        promptCacheKey: prepared.interaction.promptCacheKey,
+        tags: prepared.tags,
+    });
+}
+function tokenCount(value) {
+    return typeof value === "number" && Number.isFinite(value)
+        ? Math.max(0, Math.trunc(value))
+        : 0;
+}
+function runtimeResult(result, prepared) {
+    const decision = result.choice.decision;
+    const chain = prepared.routing.prefix.chainByModel.get(decision.selectedModel);
+    const target = prepared.routing.targetsByModel.get(decision.selectedModel);
+    const usage = result.output.usage;
+    const inputTokens = tokenCount(usage.input);
+    const cacheReadTokens = tokenCount(usage.cacheRead);
+    const cacheWriteTokens = tokenCount(usage.cacheWrite);
+    const outputTokens = tokenCount(usage.output);
+    const benchmarkReplay = estimateBenchmarkReplay({
+        enabledModels: prepared.credentials.enabled_models,
+        modelPrices: prepared.credentials.model_prices ?? {},
+        executionProvider: (model) => executionProviderForModel(prepared.credentials, model),
+        priorHit: prepared.routing.identityHit,
+        receivedAtMs: prepared.receivedAt.getTime(),
+        promptTokens: inputTokens + cacheReadTokens + cacheWriteTokens,
+        outputTokens,
+    });
+    const completedPrefix = chain === undefined || target === undefined
+        ? null
+        : buildCompletedPrefixHit({
+            request: prepared.interaction,
+            output: result.output,
+            model: decision.selectedModel,
+            nextModel: decision.selectedModel,
+            nextReasoningEffort: decision.reasoningEffort,
+            reasoningBucket: decision.reasoningEffort,
+            leaseTurnsRemaining: decision.leaseTurnsRemaining ?? null,
+            reason: decision.reason,
+            chain,
+            looseChain: prepared.routing.prefix.looseChain,
+            target,
+            conversationId: prepared.routing.conversationId,
+            leaseDepth: prepared.routing.leaseDepth,
+            providerBlockDepth: prepared.routing.providerBlockDepthFor(decision.selectedModel),
+            fingerprints: prepared.routing.prefix.fingerprints,
+            usage,
+            benchmarkReplayDepth: benchmarkReplay.depth,
+        });
+    return {
+        routing: decision,
+        output: result.output,
+        ...(completedPrefix === null
+            ? {}
+            : { prefixHit: { ...completedPrefix, updated_at: new Date().toISOString() } }),
+    };
+}
+const inProcessProtocol = {
+    responseBody: (args) => ({
+        id: args.id,
+        model: args.model,
+        output: args.output,
+        usage: args.output.usage,
+        dari_routing: {
+            selected_model: args.routing.selectedModel,
+            reasoning_effort: args.routing.reasoningEffort,
+            reason: args.routing.reason,
+            request_id: args.id,
+            ...(args.conversationId ? { conversation_id: args.conversationId } : {}),
+        },
+    }),
+    streamBody: () => {
+        throw new HttpError(400, "In-process routed runtime results only support non-streaming interactions.", "streaming_not_supported");
+    },
+};

package/dist/runtime/routing_lease.d.ts

@@ -0,0 +1,26 @@
+import { type ChatMessage, type PreviousDecision, type RoutingDecision } from "@mupt-ai/dari-router/policy-engine";
+import type { SelectorBackend } from "../strategies/routing.js";
+import { type RoutedDecision } from "../strategies/route_decision.js";
+import type { PreparedRoutedCompletion } from "./prepare_routed_completion.js";
+export type RoutingLeaseDecision = {
+    decision: RoutingDecision;
+    completedTurns: number;
+    turnInLease: number;
+};
+export declare function leasedRoutingDecision(args: {
+    messages: ChatMessage[];
+    previousDecision?: PreviousDecision;
+}): RoutingLeaseDecision | null;
+export declare function leasedDecision(prepared: PreparedRoutedCompletion): RoutedDecision | null;
+export declare function committedLeaseDecision(prepared: PreparedRoutedCompletion, decision: RoutingDecision): RoutedDecision;
+export declare function leaseFallbackDecision(args: {
+    candidates: PreparedRoutedCompletion["routing"]["candidateResolution"]["candidates"];
+    decision: RoutingDecision;
+    credentials: PreparedRoutedCompletion["credentials"];
+}): RoutingDecision | undefined;
+export declare function selectorNegotiatesLeases(backend: SelectorBackend): Promise<boolean>;
+export declare function prefetchNextLease(prepared: PreparedRoutedCompletion, selectorBackend: SelectorBackend): Promise<void>;
+export declare function leasePrefetchDue(prepared: PreparedRoutedCompletion): boolean;
+export declare function clearPendingLease(prepared: PreparedRoutedCompletion): Promise<void>;
+export declare function pendingLeaseDecision(prepared: PreparedRoutedCompletion): RoutingDecision | null;
+export declare function leaseHoldDecision(prepared: PreparedRoutedCompletion): RoutingDecision | null;