62 lines
4.2 KiB
TypeScript
62 lines
4.2 KiB
TypeScript
import assert from "node:assert/strict";
|
|
import { test } from "node:test";
|
|
import type { LanguageModelV2, LanguageModelV2CallOptions } from "@ai-sdk/provider";
|
|
import { isCredentialFailureBan, isProviderAuthenticationFailure, withAuthenticationFailover } from "@/lib/ai/runtime/authentication-failover";
|
|
|
|
const options = { prompt: [] } as unknown as LanguageModelV2CallOptions;
|
|
function model(id: string, generate: LanguageModelV2["doGenerate"], stream: LanguageModelV2["doStream"]): LanguageModelV2 {
|
|
return { specificationVersion: "v2", provider: "test-provider", modelId: id, supportedUrls: {}, doGenerate: generate, doStream: stream };
|
|
}
|
|
const authError = () => Object.assign(new Error("Invalid Authentication"), { statusCode: 401 });
|
|
const unused = async (): Promise<never> => { throw new Error("Unexpected model method"); };
|
|
|
|
test("an explicit 401 before streaming switches model once and preserves request tools", async () => {
|
|
let failures = 0; let lookups = 0; let primaryCalls = 0; let fallbackCalls = 0;
|
|
const stream = new ReadableStream({ start(controller) { controller.close(); } });
|
|
const fallback = model("working-vision", unused, async (received) => { fallbackCalls++; assert.equal(received, options); return { stream }; });
|
|
const wrapped = withAuthenticationFailover({
|
|
primary: model("rejected", unused, async () => { primaryCalls++; throw authError(); }),
|
|
onAuthenticationFailure: async () => { failures++; },
|
|
resolveFallback: async () => { lookups++; return fallback; },
|
|
});
|
|
assert.equal((await wrapped.doStream(options)).stream, stream);
|
|
await wrapped.doStream(options);
|
|
assert.equal(wrapped.modelId, "working-vision");
|
|
assert.deepEqual({ failures, lookups, primaryCalls, fallbackCalls }, { failures: 1, lookups: 1, primaryCalls: 1, fallbackCalls: 2 });
|
|
});
|
|
|
|
test("timeouts, 5xx and outcome-unknown failures never replay paid requests", async () => {
|
|
for (const error of [Object.assign(new Error("request_outcome_unknown"), { statusCode: 503 }), new Error("Request timed out"), Object.assign(new Error("Too many requests"), { statusCode: 429 })]) {
|
|
const wrapped = withAuthenticationFailover({ primary: model("primary", unused, async () => { throw error; }), onAuthenticationFailure: unused, resolveFallback: unused });
|
|
await assert.rejects(async () => await wrapped.doStream(options), (caught) => caught === error);
|
|
}
|
|
});
|
|
|
|
test("once a stream has been accepted, an error inside it is not retried", async () => {
|
|
const stream = new ReadableStream({ start(controller) { controller.enqueue({ type: "error", error: authError() }); controller.close(); } });
|
|
const wrapped = withAuthenticationFailover({ primary: model("accepted", unused, async () => ({ stream })), onAuthenticationFailure: unused, resolveFallback: unused });
|
|
assert.equal((await wrapped.doStream(options)).stream, stream);
|
|
});
|
|
|
|
test("a failing fallback does not start another retry loop", async () => {
|
|
let fallbackCalls = 0;
|
|
const fallback = model("also-rejected", unused, async () => { fallbackCalls++; throw authError(); });
|
|
const wrapped = withAuthenticationFailover({ primary: model("primary", unused, async () => { throw authError(); }), onAuthenticationFailure: async () => {}, resolveFallback: async () => fallback });
|
|
await assert.rejects(async () => await wrapped.doStream(options), /Invalid Authentication/);
|
|
assert.equal(fallbackCalls, 1);
|
|
});
|
|
|
|
test("user cancellation prevents fallback submission", async () => {
|
|
const controller = new AbortController(); controller.abort();
|
|
const wrapped = withAuthenticationFailover({ primary: model("cancelled", unused, async () => { throw authError(); }), onAuthenticationFailure: unused, resolveFallback: unused });
|
|
await assert.rejects(async () => await wrapped.doStream({ ...options, abortSignal: controller.signal }), /Invalid Authentication/);
|
|
});
|
|
|
|
test("credential failures are recognized separately from temporary outages", () => {
|
|
assert(isProviderAuthenticationFailure(new Error("Invalid Authentication")));
|
|
assert(isProviderAuthenticationFailure({ cause: { status: 401 } }));
|
|
assert(!isProviderAuthenticationFailure({ status: 503, message: "Request timed out" }));
|
|
assert(isCredentialFailureBan("provider_authentication_failed: replace API key"));
|
|
assert(!isCredentialFailureBan("request_timeout: upstream busy"));
|
|
});
|