@mupt-ai/dari-gateway 0.2.0 → 0.2.1

package/dist/providers/anthropic_safeguards.js

@@ -0,0 +1,81 @@
+import { isRecord } from "@mupt-ai/dari-router/policy-engine";
+// Claude Code's auto mode asks the Anthropic API to run its dangerous-tool-use
+// classifier server-side, at no charge, by sending `safeguards` with a paired
+// beta. The API answers with `safeguard_results`, keyed by tool_use ID, on the
+// final message_delta. Pi drops both, so the request is patched at the payload
+// boundary and the results are read off the raw SSE bytes. Other providers
+// never see the field; Claude Code then runs its own classifier for that turn.
+export function forwardsAnthropicSafeguards(request, model) {
+    return (request.anthropicSafeguards !== undefined &&
+        model.provider === "anthropic" &&
+        model.api === "anthropic-messages");
+}
+export function payloadWithSafeguards(payload, safeguards) {
+    const betas = Array.isArray(payload.betas) ? payload.betas : [];
+    return {
+        ...payload,
+        safeguards: safeguards.safeguards,
+        betas: [...new Set([...betas, ...safeguards.betas])],
+    };
+}
+export function safeguardResultsTap(fetch = globalThis.fetch) {
+    let results;
+    return {
+        results: () => results,
+        fetch: async (input, init) => {
+            const response = await fetch(input, init);
+            if (!response.ok || response.body === null)
+                return response;
+            results = undefined;
+            return new Response(response.body.pipeThrough(safeguardResultsScanner((value) => {
+                results = value;
+            })), { status: response.status, statusText: response.statusText, headers: response.headers });
+        },
+    };
+}
+export function messageWithSafeguardResults(message, results) {
+    return results === undefined ? message : { ...message, safeguardResults: results };
+}
+export async function* eventsWithSafeguardResults(events, tap) {
+    for await (const event of events) {
+        if (event.type === "done" && isRecord(event.message)) {
+            yield { ...event, message: messageWithSafeguardResults(event.message, tap.results()) };
+            continue;
+        }
+        yield event;
+    }
+}
+function safeguardResultsScanner(onResults) {
+    const decoder = new TextDecoder();
+    let pending = "";
+    const scan = (line) => {
+        if (!line.startsWith("data:") || !line.includes("\"safeguard_results\""))
+            return;
+        let event;
+        try {
+            event = JSON.parse(line.slice("data:".length));
+        }
+        catch {
+            return;
+        }
+        if (isRecord(event) &&
+            event.type === "message_delta" &&
+            isRecord(event.delta) &&
+            Array.isArray(event.delta.safeguard_results)) {
+            onResults(event.delta.safeguard_results);
+        }
+    };
+    return new TransformStream({
+        transform(chunk, controller) {
+            controller.enqueue(chunk);
+            pending += decoder.decode(chunk, { stream: true });
+            const lines = pending.split("\n");
+            pending = lines.pop() ?? "";
+            for (const line of lines)
+                scan(line.trimEnd());
+        },
+        flush() {
+            scan((pending + decoder.decode()).trimEnd());
+        },
+    });
+}

package/dist/providers/aws_assume_role.d.ts

@@ -0,0 +1,17 @@
+import type { AwsAssumeRoleProviderAuth, AwsSigV4ProviderAuth, ProviderAuth, ResolvedProviderAuth } from "../types.js";
+export type AssumedRoleCredentials = AwsSigV4ProviderAuth & {
+    session_token: string;
+    expires_at: number;
+};
+export type AssumeRoleFn = (auth: AwsAssumeRoleProviderAuth) => Promise<AssumedRoleCredentials>;
+export type AwsRoleCredentials = {
+    /** Swap an AssumeRole auth for temporary SigV4 credentials; other auth passes through. */
+    resolve(auth: ProviderAuth): Promise<ResolvedProviderAuth>;
+    /** Temporary credential values already issued for this role, for log scrubbing. */
+    secrets(auth: AwsAssumeRoleProviderAuth): string[];
+};
+export declare function createAwsRoleCredentials(deps: {
+    assumeRole: AssumeRoleFn;
+    now?: () => number;
+}): AwsRoleCredentials;
+export declare const awsRoleCredentials: AwsRoleCredentials;

package/dist/providers/aws_assume_role.js

@@ -0,0 +1,107 @@
+import { AssumeRoleCommand, STSClient } from "@aws-sdk/client-sts";
+import { HttpError } from "../http/errors.js";
+const ROLE_SESSION_NAME = "dari-routing";
+// AssumeRole always goes through this regional STS endpoint. The customer's
+// Bedrock region may be an opt-in region Dari's own account has not enabled,
+// where a regional STS call would fail. Regional endpoints issue v2 tokens
+// that are valid in every region, so the temporary credentials still sign
+// Bedrock requests in the customer's region.
+const STS_REGION = "us-east-1";
+// Re-assume the role this long before the temporary credentials expire so a
+// long Bedrock stream never starts on credentials that lapse mid-request.
+const REFRESH_MARGIN_MS = 5 * 60_000;
+export function createAwsRoleCredentials(deps) {
+    const now = deps.now ?? Date.now;
+    const issued = new Map();
+    // Credentials replaced by a refresh stay scrubbable until they expire:
+    // a stream that started before the refresh still signs with them.
+    const retired = new Map();
+    const inflight = new Map();
+    function assume(auth) {
+        const key = cacheKey(auth);
+        const cached = issued.get(key);
+        if (cached && cached.expires_at - now() > REFRESH_MARGIN_MS) {
+            return Promise.resolve(cached);
+        }
+        let pending = inflight.get(key);
+        if (!pending) {
+            pending = deps.assumeRole(auth)
+                .then((credentials) => {
+                const previous = issued.get(key);
+                if (previous) {
+                    retired.set(key, [
+                        ...(retired.get(key) ?? []).filter((c) => c.expires_at > now()),
+                        previous,
+                    ]);
+                }
+                issued.set(key, credentials);
+                return credentials;
+            })
+                .finally(() => inflight.delete(key));
+            inflight.set(key, pending);
+        }
+        return pending;
+    }
+    return {
+        async resolve(auth) {
+            if (auth.type !== "aws_assume_role")
+                return auth;
+            const { expires_at: _expiresAt, ...credentials } = await assume(auth);
+            return credentials;
+        },
+        secrets(auth) {
+            const key = cacheKey(auth);
+            const current = issued.get(key);
+            const live = [
+                ...(current ? [current] : []),
+                ...(retired.get(key) ?? []).filter((c) => c.expires_at > now()),
+            ];
+            return live.flatMap((c) => [c.access_key_id, c.secret_access_key, c.session_token]);
+        },
+    };
+}
+function cacheKey(auth) {
+    return `${auth.role_arn}\n${auth.external_id}\n${auth.region}`;
+}
+async function stsAssumeRole(auth) {
+    const accessKeyId = process.env.DARI_AWS_ASSUME_ROLE_ACCESS_KEY_ID?.trim();
+    const secretAccessKey = process.env.DARI_AWS_ASSUME_ROLE_SECRET_ACCESS_KEY?.trim();
+    if (!accessKeyId || !secretAccessKey) {
+        throw new HttpError(500, "DARI_AWS_ASSUME_ROLE_ACCESS_KEY_ID and DARI_AWS_ASSUME_ROLE_SECRET_ACCESS_KEY are not configured.", "configuration_error");
+    }
+    const client = new STSClient({
+        region: STS_REGION,
+        credentials: { accessKeyId, secretAccessKey },
+    });
+    let result;
+    try {
+        result = await client.send(new AssumeRoleCommand({
+            RoleArn: auth.role_arn,
+            ExternalId: auth.external_id,
+            RoleSessionName: ROLE_SESSION_NAME,
+        }));
+    }
+    catch (error) {
+        const reason = error instanceof Error ? error.message : String(error);
+        throw new Error(`AWS AssumeRole for ${auth.role_arn} failed: ${reason}`);
+    }
+    finally {
+        client.destroy();
+    }
+    const credentials = result.Credentials;
+    if (!credentials?.AccessKeyId
+        || !credentials.SecretAccessKey
+        || !credentials.SessionToken
+        || !credentials.Expiration) {
+        throw new Error(`AWS AssumeRole for ${auth.role_arn} returned incomplete credentials.`);
+    }
+    return {
+        type: "aws_sigv4",
+        access_key_id: credentials.AccessKeyId,
+        secret_access_key: credentials.SecretAccessKey,
+        session_token: credentials.SessionToken,
+        region: auth.region,
+        expires_at: credentials.Expiration.getTime(),
+    };
+}
+export const awsRoleCredentials = createAwsRoleCredentials({ assumeRole: stsAssumeRole });

package/dist/providers/cancel_estimate.d.ts

@@ -0,0 +1,43 @@
+import type { CandidateCostEstimate, ModelPricing, ReasoningEffort } from "@mupt-ai/dari-router/policy-engine";
+import type { InteractionOutput, InteractionUsage, ProviderIdentity } from "../interaction/types.js";
+/** Prose from current models averages about 4.8 characters per token. */
+export declare const ESTIMATED_CHARS_PER_TOKEN = 4.8;
+export type PromptTokenEstimate = {
+    promptTokens: number;
+    warmTokens: number;
+};
+export type DeliveredOutput = {
+    text: string;
+    textChars: number;
+    reasoningChars: number;
+};
+export declare const NO_DELIVERED_OUTPUT: DeliveredOutput;
+export declare function promptTokenEstimateFor(estimates: readonly CandidateCostEstimate[] | undefined, model: string, reasoningEffort: ReasoningEffort): PromptTokenEstimate | undefined;
+/** Measure the provider text the protocol actually consumed for the client. */
+export declare function observeDeliveredOutput(events: AsyncIterable<Record<string, unknown>>): {
+    events: AsyncIterable<Record<string, unknown>>;
+    delivered: () => DeliveredOutput;
+};
+/**
+ * Usage for a cancelled attempt: prompt tokens from the selector's estimate,
+ * output tokens from the characters delivered before the abort. Reasoning the
+ * provider never surfaced is not counted, so the estimate errs in the
+ * customer's favor. Cost is zero when the model has no known pricing.
+ */
+export declare function estimatedCancelUsage(args: {
+    pricing: ModelPricing | null;
+    promptEstimate: PromptTokenEstimate | undefined;
+    delivered: DeliveredOutput;
+}): InteractionUsage & {
+    estimated: true;
+};
+/**
+ * The delivered partial as an assistant turn, so the interrupted request can
+ * still write a prefix entry that the follow-up turn may match.
+ */
+export declare function partialInteractionOutput(args: {
+    delivered: DeliveredOutput;
+    usage: InteractionUsage;
+    source: ProviderIdentity;
+    createdAtMs: number;
+}): InteractionOutput | null;

package/dist/providers/cancel_estimate.js

@@ -0,0 +1,89 @@
+import { pricingForPrompt } from "@mupt-ai/dari-router/policy-engine";
+/** Prose from current models averages about 4.8 characters per token. */
+export const ESTIMATED_CHARS_PER_TOKEN = 4.8;
+const MTOK = 1_000_000;
+const MAX_RETAINED_TEXT_CHARS = 512 * 1024;
+export const NO_DELIVERED_OUTPUT = {
+    text: "",
+    textChars: 0,
+    reasoningChars: 0,
+};
+export function promptTokenEstimateFor(estimates, model, reasoningEffort) {
+    const match = estimates?.find((item) => item.model === model && item.reasoning_effort === reasoningEffort) ?? estimates?.find((item) => item.model === model);
+    return match
+        ? { promptTokens: match.est_prompt_tokens, warmTokens: match.warm_tokens }
+        : undefined;
+}
+/** Measure the provider text the protocol actually consumed for the client. */
+export function observeDeliveredOutput(events) {
+    const state = { text: "", textChars: 0, reasoningChars: 0 };
+    async function* observed() {
+        for await (const event of events) {
+            if (typeof event.delta === "string") {
+                if (event.type === "text_delta") {
+                    state.textChars += event.delta.length;
+                    if (state.text.length < MAX_RETAINED_TEXT_CHARS) {
+                        state.text += event.delta.slice(0, MAX_RETAINED_TEXT_CHARS - state.text.length);
+                    }
+                }
+                else if (event.type === "thinking_delta") {
+                    state.reasoningChars += event.delta.length;
+                }
+            }
+            yield event;
+        }
+    }
+    return { events: observed(), delivered: () => ({ ...state }) };
+}
+/**
+ * Usage for a cancelled attempt: prompt tokens from the selector's estimate,
+ * output tokens from the characters delivered before the abort. Reasoning the
+ * provider never surfaced is not counted, so the estimate errs in the
+ * customer's favor. Cost is zero when the model has no known pricing.
+ */
+export function estimatedCancelUsage(args) {
+    const cacheRead = Math.max(0, Math.round(args.promptEstimate?.warmTokens ?? 0));
+    const input = Math.max(0, Math.round((args.promptEstimate?.promptTokens ?? 0) - cacheRead));
+    const output = Math.round((args.delivered.textChars + args.delivered.reasoningChars) /
+        ESTIMATED_CHARS_PER_TOKEN);
+    const price = pricingForPrompt(args.pricing ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, input + cacheRead);
+    const cost = {
+        input: (input * price.input) / MTOK,
+        output: (output * price.output) / MTOK,
+        cacheRead: (cacheRead * price.cacheRead) / MTOK,
+        cacheWrite: 0,
+        total: 0,
+    };
+    cost.total = cost.input + cost.output + cost.cacheRead;
+    return {
+        input,
+        output,
+        cacheRead,
+        cacheWrite: 0,
+        totalTokens: input + cacheRead + output,
+        cost,
+        estimated: true,
+    };
+}
+/**
+ * The delivered partial as an assistant turn, so the interrupted request can
+ * still write a prefix entry that the follow-up turn may match.
+ */
+export function partialInteractionOutput(args) {
+    if (args.delivered.text.length === 0)
+        return null;
+    return {
+        createdAtMs: args.createdAtMs,
+        source: args.source,
+        items: [
+            {
+                type: "message",
+                role: "assistant",
+                content: [{ type: "text", text: args.delivered.text }],
+                contentShape: "parts",
+            },
+        ],
+        usage: args.usage,
+        stopReason: "aborted",
+    };
+}

package/dist/providers/execution.d.ts

@@ -0,0 +1,72 @@
+import type { RoutingDecision } from "@mupt-ai/dari-router/policy-engine";
+import { piMessageToInteractionOutput } from "./pi_output.js";
+import type { InteractionRequest } from "../interaction/types.js";
+import { type RoutedRuntimeReporter } from "../runtime/runtime_reporter.js";
+import type { RoutedProtocol } from "../runtime/routed_protocol.js";
+import { type ProviderExecution, type StartedCompletion, type StartedStream } from "./provider_call.js";
+export type ProviderChoice = {
+    decision: RoutingDecision;
+    execution: ProviderExecution;
+    role: "primary" | "fallback";
+};
+export type ResolveFallbackChoice = () => Promise<ProviderChoice | null>;
+export type PreparedStreamExecution = {
+    choice: ProviderChoice;
+    started: StartedStream;
+    attempt: number;
+    failedSetupLatencyMs: number;
+};
+export type CompletedNonStreamingExecution = {
+    choice: ProviderChoice;
+    output: ReturnType<typeof piMessageToInteractionOutput>;
+    responseBody: Record<string, unknown>;
+    reportBarrierError?: unknown;
+};
+export declare function executeNonStreaming(args: {
+    requestId: string;
+    interaction: InteractionRequest;
+    primary: ProviderChoice;
+    primaryRetries: number;
+    resolveFallback: ResolveFallbackChoice;
+    reporter: RoutedRuntimeReporter;
+    protocol: RoutedProtocol;
+    startedCompletion?: StartedCompletion;
+    reportBarrier?: Promise<unknown>;
+    clientSignal?: AbortSignal;
+}): Promise<Response>;
+export declare function completeNonStreaming(args: {
+    requestId: string;
+    interaction: InteractionRequest;
+    primary: ProviderChoice;
+    primaryRetries: number;
+    resolveFallback: ResolveFallbackChoice;
+    reporter: RoutedRuntimeReporter;
+    protocol: RoutedProtocol;
+    startedCompletion?: StartedCompletion;
+    reportBarrier?: Promise<unknown>;
+    clientSignal?: AbortSignal;
+}): Promise<CompletedNonStreamingExecution>;
+export declare function prepareStreamingExecution(args: {
+    requestId: string;
+    interaction: InteractionRequest;
+    primary: ProviderChoice;
+    primaryRetries: number;
+    resolveFallback: ResolveFallbackChoice;
+    reporter: RoutedRuntimeReporter;
+    startedStream?: StartedStream;
+    initialError?: unknown;
+    reportBarrier?: Promise<unknown>;
+    clientSignal?: AbortSignal;
+}): Promise<PreparedStreamExecution>;
+export declare function executeStreaming(args: {
+    requestId: string;
+    interaction: InteractionRequest;
+    choice: ProviderChoice;
+    reporter: RoutedRuntimeReporter;
+    protocol: RoutedProtocol;
+    startedStream: StartedStream;
+    attempt: number;
+    failedSetupLatencyMs: number;
+    reportBarrier?: Promise<unknown>;
+    clientSignal?: AbortSignal;
+}): Promise<Response>;