@dbx-tools/appkit-model-gateway 0.9.201 → 0.9.203

package/src/gateway.ts

@@ -43,6 +43,7 @@ import {
 
 import { AiDecideRouterService } from "./ai-decide.ts";
 import { listModelsPayload } from "./models.ts";
+import { ProfileFallback, type ModelGatewayProfile } from "./profile-fallback.ts";
 import { adaptInferenceReasoning, learnAndAdaptReasoningRetry } from "./reasoning-adapt.ts";
 import {
   DatabricksModelRegistry,
@@ -78,6 +79,7 @@ export interface ModelGatewayOptions {
   readonly model?: string;
   readonly modelClass?: ModelClassType;
   readonly preferSystemModels?: boolean;
+  readonly profiles?: readonly ModelGatewayProfile[];
 }
 
 /** Error carrying an HTTP status and protocol-safe message. */
@@ -100,14 +102,22 @@ export class ModelGateway {
   private readonly model?: string;
   private readonly modelClass?: ModelClassType;
   private readonly preferSystemModels: boolean;
+  private readonly profileFallback?: ProfileFallback;
+  private readonly primaryProfile?: ModelGatewayProfile;
 
   constructor(options: ModelGatewayOptions = {}, registry?: ModelRegistry) {
+    this.primaryProfile = options.profiles?.[0];
+    this.profileFallback =
+      options.profiles && options.profiles.length > 1
+        ? new ProfileFallback(options.profiles)
+        : undefined;
     const registryOptions: ModelRegistryOptions = {
       ...(options.cacheTtlMs !== undefined ? { ttlMs: options.cacheTtlMs } : {}),
       ...(options.overrides ? { overrides: options.overrides } : {}),
       ...(options.preferSystemModels !== undefined
         ? { preferSystemModels: options.preferSystemModels }
         : {}),
+      ...(this.primaryProfile ? { client: this.primaryProfile.client } : {}),
     };
     this.registry = registry ?? new DatabricksModelRegistry(registryOptions);
     this.aiDecide = options.aiDecideService ?? new AiDecideRouterService();
@@ -129,6 +139,7 @@ export class ModelGateway {
         },
         options.smartRouterService,
         options.smartRouterStore,
+        this.primaryProfile?.client,
       );
     }
     this.model = options.model;
@@ -174,6 +185,7 @@ export class ModelGateway {
         targets
           .map((target) => target.endpoint)
           .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined),
+        this.primaryProfile?.client,
       ));
     if (this.smartRouterMode === "true" && !smartRouterEnabled) {
       throw new ModelGatewayError(
@@ -251,7 +263,7 @@ export class ModelGateway {
           (registryTargets ?? [])
             .map((candidate) => candidate.endpoint)
             .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined),
-          undefined,
+          this.primaryProfile?.client,
           {
             strategy: aiDecideStrategy(headers),
             filters: {
@@ -344,9 +356,25 @@ export class ModelGateway {
     signal: AbortSignal,
   ): Promise<Response> {
     if (route.upstreamProtocol === "ai-sdk") {
-      return translateGatewayRequest(route, body, signal);
+      return translateGatewayRequest(route, body, signal, {
+        ...(this.primaryProfile ? { client: this.primaryProfile.client } : {}),
+        ...(this.profileFallback
+          ? { fetch: this.profileFallback.createFetch(route.upstreamModel) }
+          : {}),
+      });
+    }
+    if (this.profileFallback) {
+      return this.profileFallback.request(route.upstreamModel, signal, ({ client }) =>
+        fetchDatabricks({ route, body, headers, signal, client }),
+      );
     }
-    return fetchDatabricks({ route, body, headers, signal });
+    return fetchDatabricks({
+      route,
+      body,
+      headers,
+      signal,
+      ...(this.primaryProfile ? { client: this.primaryProfile.client } : {}),
+    });
   }
 }
 

package/src/profile-fallback.ts

@@ -0,0 +1,277 @@
+/**
+ * Ordered Databricks profile fallback for rate-limited model requests.
+ *
+ * @module
+ */
+
+import type { WorkspaceClient } from "@databricks/appkit";
+import { log } from "@dbx-tools/shared-core";
+
+import { fetchWorkspace } from "./transport.ts";
+
+const logger = log.logger("appkit/model-gateway/profile-fallback");
+const GENERIC_RATE_LIMIT_COOLDOWN_MS = 15_000;
+const WORKSPACE_TOKEN_COOLDOWN_MS = 60_000;
+
+/** One configured workspace identity in fallback order. */
+export interface ModelGatewayProfile {
+  readonly profile?: string;
+  readonly workspace: string;
+  readonly client: WorkspaceClient;
+}
+
+/** Parsed upstream rate-limit timing and diagnostic fields. */
+export interface ModelGatewayRateLimit {
+  readonly retryAfterMs: number;
+  readonly source: "header" | "body" | "workspace-token-window" | "default";
+  readonly errorCode?: string;
+  readonly message?: string;
+}
+
+/** Injectable time and transport boundaries for profile-fallback tests. */
+export interface ProfileFallbackDependencies {
+  readonly now?: () => number;
+  readonly sleep?: (delayMs: number, signal: AbortSignal) => Promise<void>;
+  readonly fetch?: typeof globalThis.fetch;
+}
+
+/** Process-local ordered profile cooldowns keyed by workspace and model. */
+export class ProfileFallback {
+  private readonly cooldowns = new Map<string, number>();
+  private readonly now: () => number;
+  private readonly sleep: (delayMs: number, signal: AbortSignal) => Promise<void>;
+  private readonly fetcher: typeof globalThis.fetch;
+
+  constructor(
+    readonly profiles: readonly ModelGatewayProfile[],
+    dependencies: ProfileFallbackDependencies = {},
+  ) {
+    if (profiles.length < 2) {
+      throw new TypeError("Profile fallback requires at least two profiles");
+    }
+    this.now = dependencies.now ?? Date.now;
+    this.sleep = dependencies.sleep ?? abortableSleep;
+    this.fetcher = dependencies.fetch ?? globalThis.fetch.bind(globalThis);
+  }
+
+  /** Primary profile used for discovery and workspace-scoped routing. */
+  get primary(): ModelGatewayProfile {
+    return this.profiles[0]!;
+  }
+
+  /** Execute against the first available profile and cool down HTTP 429 responders. */
+  async request(
+    model: string,
+    signal: AbortSignal,
+    send: (profile: ModelGatewayProfile) => Promise<Response>,
+  ): Promise<Response> {
+    const attempted = new Set<number>();
+    let waitedForExistingCooldown = false;
+    while (true) {
+      const now = this.now();
+      const profileIndex = this.availableProfile(model, attempted, now);
+      if (profileIndex === undefined) {
+        if (attempted.size > 0 || waitedForExistingCooldown) {
+          throw new Error("Profile fallback exhausted without an upstream response");
+        }
+        const delayMs = this.nextCooldownDelay(model, now);
+        if (delayMs === undefined) {
+          throw new Error("Profile fallback has no available profile");
+        }
+        waitedForExistingCooldown = true;
+        await this.sleep(delayMs, signal);
+        continue;
+      }
+
+      attempted.add(profileIndex);
+      const profile = this.profiles[profileIndex]!;
+      const response = await send(profile);
+      if (response.status !== 429) {
+        this.cooldowns.delete(this.key(profile.workspace, model));
+        if (profileIndex > 0) {
+          logger.info("model request completed through fallback profile", {
+            model,
+            profile: profile.profile ?? "default",
+            profileIndex,
+          });
+        }
+        return response;
+      }
+
+      const rateLimit = await modelGatewayRateLimit(response, this.now());
+      this.cooldowns.set(this.key(profile.workspace, model), this.now() + rateLimit.retryAfterMs);
+      logger.warn("model profile rate limited", {
+        errorCode: rateLimit.errorCode,
+        upstreamMessage: rateLimit.message,
+        model,
+        profile: profile.profile ?? "default",
+        profileIndex,
+        retryAfterMs: rateLimit.retryAfterMs,
+        retryAfterSource: rateLimit.source,
+        workspace: profile.workspace,
+      });
+
+      if (this.availableProfile(model, attempted, this.now()) === undefined) {
+        return response;
+      }
+      await response.body?.cancel().catch(() => undefined);
+    }
+  }
+
+  /** Build an AI SDK fetch that replays a request through the ordered profile pool. */
+  createFetch(model: string): typeof globalThis.fetch {
+    return (async (input, init) => {
+      const request = new Request(input, init);
+      const body =
+        request.method === "GET" || request.method === "HEAD"
+          ? undefined
+          : await request.arrayBuffer();
+      const originalUrl = new URL(request.url);
+      const originalHeaders = new Headers(request.headers);
+      return this.request(model, request.signal, async ({ client, workspace }) => {
+        const url = new URL(`${originalUrl.pathname}${originalUrl.search}`, workspace);
+        return fetchWorkspace(
+          client,
+          url,
+          {
+            method: request.method,
+            headers: originalHeaders,
+            body,
+            signal: request.signal,
+          },
+          this.fetcher,
+        );
+      });
+    }) as typeof globalThis.fetch;
+  }
+
+  private availableProfile(
+    model: string,
+    attempted: ReadonlySet<number>,
+    now: number,
+  ): number | undefined {
+    for (let index = 0; index < this.profiles.length; index += 1) {
+      if (attempted.has(index)) continue;
+      const workspace = this.profiles[index]!.workspace;
+      if ((this.cooldowns.get(this.key(workspace, model)) ?? 0) <= now) return index;
+    }
+    return undefined;
+  }
+
+  private nextCooldownDelay(model: string, now: number): number | undefined {
+    const cooldowns = this.profiles
+      .map((profile) => this.cooldowns.get(this.key(profile.workspace, model)))
+      .filter((until): until is number => until !== undefined && until > now);
+    if (cooldowns.length === 0) return undefined;
+    return Math.max(0, Math.min(...cooldowns) - now);
+  }
+
+  private key(workspace: string, model: string): string {
+    return `${normalizeWorkspace(workspace)}\0${model}`;
+  }
+}
+
+function normalizeWorkspace(workspace: string): string {
+  const url = new URL(workspace);
+  return url.origin.toLowerCase();
+}
+
+/** Parse HTTP and Databricks JSON retry timing for a 429 response. */
+export async function modelGatewayRateLimit(
+  response: Response,
+  now = Date.now(),
+): Promise<ModelGatewayRateLimit> {
+  const headerDelay = retryAfterHeaderMs(response.headers.get("retry-after"), now);
+  const body = await response
+    .clone()
+    .json()
+    .catch(() => undefined);
+  const details = rateLimitBody(body);
+  if (headerDelay !== undefined) {
+    return { ...details, retryAfterMs: headerDelay, source: "header" };
+  }
+  if (details.retryAfterMs !== undefined) {
+    return { ...details, retryAfterMs: details.retryAfterMs, source: "body" };
+  }
+  if (isWorkspaceTokenRateLimit(details.message)) {
+    return {
+      ...details,
+      retryAfterMs: WORKSPACE_TOKEN_COOLDOWN_MS,
+      source: "workspace-token-window",
+    };
+  }
+  return {
+    ...details,
+    retryAfterMs: GENERIC_RATE_LIMIT_COOLDOWN_MS,
+    source: "default",
+  };
+}
+
+function rateLimitBody(value: unknown): {
+  readonly retryAfterMs?: number;
+  readonly errorCode?: string;
+  readonly message?: string;
+} {
+  const root = record(value);
+  const error = record(root?.error) ?? root;
+  const retryAfter = integer(error?.retry_after ?? root?.retry_after);
+  const errorCode = text(error?.error_code ?? root?.error_code);
+  const message =
+    text(error?.message ?? root?.message) ??
+    (typeof root?.error === "string" ? text(root.error) : undefined);
+  return {
+    ...(retryAfter !== undefined ? { retryAfterMs: retryAfter * 1_000 } : {}),
+    ...(errorCode ? { errorCode } : {}),
+    ...(message ? { message } : {}),
+  };
+}
+
+function retryAfterHeaderMs(value: string | null, now: number): number | undefined {
+  const normalized = value?.trim();
+  if (!normalized) return undefined;
+  const seconds = Number(normalized);
+  if (Number.isFinite(seconds) && seconds >= 0) return seconds * 1_000;
+  const timestamp = Date.parse(normalized);
+  return Number.isFinite(timestamp) ? Math.max(0, timestamp - now) : undefined;
+}
+
+function isWorkspaceTokenRateLimit(message: string | undefined): boolean {
+  return (
+    message !== undefined && /exceeded workspace .*tokens per minute rate limit/i.test(message)
+  );
+}
+
+function integer(value: unknown): number | undefined {
+  const normalized = typeof value === "string" ? value.trim().replaceAll(",", "") : value;
+  const parsed = typeof normalized === "number" ? normalized : Number(normalized);
+  return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : undefined;
+}
+
+function text(value: unknown): string | undefined {
+  return typeof value === "string" && value.trim() ? value.trim() : undefined;
+}
+
+function record(value: unknown): Record<string, unknown> | undefined {
+  return value !== null && typeof value === "object" && !Array.isArray(value)
+    ? (value as Record<string, unknown>)
+    : undefined;
+}
+
+function abortableSleep(delayMs: number, signal: AbortSignal): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (signal.aborted) {
+      reject(signal.reason);
+      return;
+    }
+    const timeout = setTimeout(finish, delayMs);
+    function finish() {
+      signal.removeEventListener("abort", abort);
+      resolve();
+    }
+    function abort() {
+      clearTimeout(timeout);
+      reject(signal.reason);
+    }
+    signal.addEventListener("abort", abort, { once: true });
+  });
+}

package/src/translation.ts

@@ -4,7 +4,7 @@
  * @module
  */
 
-import { getExecutionContext } from "@databricks/appkit";
+import { getExecutionContext, type WorkspaceClient } from "@databricks/appkit";
 import { invoke } from "@dbx-tools/model";
 import {
   decodeGatewayRequest,
@@ -26,6 +26,10 @@ export async function translateGatewayRequest(
   route: GatewayRoute,
   body: Readonly<Record<string, unknown>>,
   signal: AbortSignal,
+  options: {
+    readonly client?: WorkspaceClient;
+    readonly fetch?: typeof globalThis.fetch;
+  } = {},
 ): Promise<Response> {
   const decoded = decodeGatewayRequest(route.clientProtocol, body);
   logger.debug("decoded request", {
@@ -35,7 +39,7 @@ export async function translateGatewayRequest(
     stream: body.stream === true,
     toolCount: Object.keys(decoded.tools).length,
   });
-  const model = await translationModel(route);
+  const model = await translationModel(route, options.client, options.fetch);
   const result = streamText({
     model,
     messages: decoded.messages,
@@ -110,8 +114,11 @@ export async function translateGatewayRequest(
   );
 }
 
-async function translationModel(route: GatewayRoute): Promise<LanguageModel> {
-  const client = getExecutionContext().client;
+async function translationModel(
+  route: GatewayRoute,
+  client: WorkspaceClient = getExecutionContext().client,
+  fetcher?: typeof globalThis.fetch,
+): Promise<LanguageModel> {
   const host = (await client.config.getHost()).toString();
   const headers = await invoke.authHeaders(client);
   const protocol = translationProtocol(route);
@@ -121,6 +128,7 @@ async function translationModel(route: GatewayRoute): Promise<LanguageModel> {
     protocol,
     host,
     headers,
+    ...(fetcher ? { fetch: fetcher } : {}),
   });
 }
 

package/src/transport.ts

@@ -4,7 +4,7 @@
  * @module
  */
 
-import { getExecutionContext } from "@databricks/appkit";
+import { getExecutionContext, type WorkspaceClient } from "@databricks/appkit";
 import { refreshWorkspaceClientAuthentication } from "@dbx-tools/databricks";
 import { log, toError } from "@dbx-tools/shared-core";
 import { stripUnsupportedChatFields } from "@dbx-tools/shared-model";
@@ -45,11 +45,12 @@ export interface DatabricksTransportRequest {
   readonly body: Readonly<Record<string, unknown>>;
   readonly headers: Headers;
   readonly signal: AbortSignal;
+  readonly client?: WorkspaceClient;
 }
 
 /** Execute a direct route with SDK-owned authentication and an unbuffered response body. */
 export async function fetchDatabricks(request: DatabricksTransportRequest): Promise<Response> {
-  const client = getExecutionContext().client;
+  const client = request.client ?? getExecutionContext().client;
   const host = (await client.config.getHost()).toString();
   const headers = new Headers({ "content-type": "application/json" });
   request.headers.forEach((value, name) => {
@@ -58,7 +59,6 @@ export async function fetchDatabricks(request: DatabricksTransportRequest): Prom
       headers.set(name, value);
     }
   });
-  await client.config.authenticate(headers);
   const url = upstreamUrl(host, request.route);
   const startedAt = Date.now();
   const body = JSON.stringify(directDatabricksRequestBody(request.route, request.body));
@@ -69,28 +69,12 @@ export async function fetchDatabricks(request: DatabricksTransportRequest): Prom
     upstreamProtocol: request.route.upstreamProtocol,
   });
   try {
-    let response = await fetch(url, {
+    const response = await fetchWorkspace(client, url, {
       method: "POST",
       headers,
       body,
       signal: request.signal,
     });
-    if (response.status === 401 && (await refreshWorkspaceClientAuthentication(client))) {
-      logger.debug("retrying rejected upstream credential", {
-        model: request.route.upstreamModel,
-        upstreamProtocol: request.route.upstreamProtocol,
-      });
-      await response.body?.cancel();
-      headers.delete("authorization");
-      headers.delete("x-databricks-workspace-id");
-      await client.config.authenticate(headers);
-      response = await fetch(url, {
-        method: "POST",
-        headers,
-        body,
-        signal: request.signal,
-      });
-    }
     const details = {
       elapsedMs: Date.now() - startedAt,
       model: request.route.upstreamModel,
@@ -117,6 +101,32 @@ export async function fetchDatabricks(request: DatabricksTransportRequest): Prom
   }
 }
 
+/** Execute one authenticated workspace request and refresh a rejected credential once. */
+export async function fetchWorkspace(
+  client: WorkspaceClient,
+  input: string | URL,
+  init: RequestInit,
+  fetcher: typeof globalThis.fetch = globalThis.fetch.bind(globalThis),
+): Promise<Response> {
+  const headers = new Headers(init.headers);
+  headers.delete("authorization");
+  headers.delete("x-databricks-workspace-id");
+  await client.config.authenticate(headers);
+  let response = await fetcher(input, { ...init, headers });
+  if (response.status !== 401 || !(await refreshWorkspaceClientAuthentication(client))) {
+    return response;
+  }
+  logger.debug("retrying rejected workspace credential", {
+    path: new URL(input).pathname,
+  });
+  await response.body?.cancel();
+  headers.delete("authorization");
+  headers.delete("x-databricks-workspace-id");
+  await client.config.authenticate(headers);
+  response = await fetcher(input, { ...init, headers });
+  return response;
+}
+
 /** Remove client compatibility fields rejected by direct Databricks protocols. */
 export function directDatabricksRequestBody(
   route: Pick<GatewayRoute, "upstreamProtocol" | "upstreamModel">,