diff --git a/packages/server/src/routes/connections.routes.ts b/packages/server/src/routes/connections.routes.ts index 979eaf94cb..81a489e1e0 100644 --- a/packages/server/src/routes/connections.routes.ts +++ b/packages/server/src/routes/connections.routes.ts @@ -460,6 +460,7 @@ export async function connectionsRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); let responseText = ""; for await (const chunk of provider.chat([{ role: "user", content: "Reply with OK." }], { @@ -502,6 +503,7 @@ export async function connectionsRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); let responseText = ""; for await (const chunk of provider.chat([{ role: "user", content: "Reply with OK." }], { @@ -1369,6 +1371,7 @@ export async function connectionsRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); let fullResponse = ""; diff --git a/packages/server/src/routes/encounter.routes.ts b/packages/server/src/routes/encounter.routes.ts index 1d0a9f670b..b1862ed3f3 100644 --- a/packages/server/src/routes/encounter.routes.ts +++ b/packages/server/src/routes/encounter.routes.ts @@ -589,6 +589,7 @@ export async function encounterRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); const characterIds: string[] = JSON.parse(chat.characterIds as string); @@ -699,6 +700,7 @@ export async function encounterRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); const characterIds: string[] = JSON.parse(chat.characterIds as string); @@ -800,6 +802,7 @@ export async function encounterRoutes(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); const characterIds: string[] = JSON.parse(chat.characterIds as string); diff --git a/packages/server/src/routes/generate/dry-run-route.ts b/packages/server/src/routes/generate/dry-run-route.ts index 337f1711b2..fc8b66c042 100644 --- a/packages/server/src/routes/generate/dry-run-route.ts +++ b/packages/server/src/routes/generate/dry-run-route.ts @@ -1615,6 +1615,7 @@ export async function registerDryRunRoute(app: FastifyInstance) { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + connId ?? undefined, ); // ── Mirror /api/generate: normalize + fit prompt to context ── diff --git a/packages/server/src/routes/lorebooks.routes.ts b/packages/server/src/routes/lorebooks.routes.ts index 04fd4178ef..249fbf91ad 100644 --- a/packages/server/src/routes/lorebooks.routes.ts +++ b/packages/server/src/routes/lorebooks.routes.ts @@ -1193,6 +1193,7 @@ export async function lorebooksRoutes(app: FastifyInstance) { resolvedConn.claudeFastMode === "true", resolvedConn.treatAsLocalEndpoint === "true", resolvedConn.defaultParameters, + resolvedConn.id, ); })(); const embeddingModel = useLocalSidecar ? LOCAL_SIDECAR_MODEL : body.model; diff --git a/packages/server/src/services/capability-packages/capability-language-model.service.ts b/packages/server/src/services/capability-packages/capability-language-model.service.ts index 5f9f45f96c..26270ffc8b 100644 --- a/packages/server/src/services/capability-packages/capability-language-model.service.ts +++ b/packages/server/src/services/capability-packages/capability-language-model.service.ts @@ -10,6 +10,7 @@ import type { DB } from "../../db/connection.js"; import { fitMessagesToContext, type BaseLLMProvider, type ChatMessage, type ChatOptions } from "../llm/base-provider.js"; import { getLocalSidecarProvider, LOCAL_SIDECAR_MODEL } from "../llm/local-sidecar.js"; import { createLLMProvider } from "../llm/provider-registry.js"; +import { unwrapConnectionAdmissionProvider } from "../generation/connection-admission.js"; import { createConnectionsStorage } from "../storage/connections.storage.js"; export function createCapabilityLanguageModelHost(db: DB): CapabilityLanguageModelHost { @@ -21,7 +22,7 @@ export function createCapabilityLanguageModelHost(db: DB): CapabilityLanguageMod }; const resolvedModel = (provider: BaseLLMProvider, connectionId: string, model: string) => Object.freeze({ - name: provider.constructor.name, + name: unwrapConnectionAdmissionProvider(provider).constructor.name, connectionId, model, maxContext: provider.maxContextValue, @@ -77,6 +78,8 @@ export function createCapabilityLanguageModelHost(db: DB): CapabilityLanguageMod connection.maxTokensOverride, connection.claudeFastMode === "true", connection.treatAsLocalEndpoint === "true", + undefined, + connection.id, ), connection.id, requireModel(model ?? connection.model), diff --git a/packages/server/src/services/generation/connection-admission.ts b/packages/server/src/services/generation/connection-admission.ts new file mode 100644 index 0000000000..6c36f6fff2 --- /dev/null +++ b/packages/server/src/services/generation/connection-admission.ts @@ -0,0 +1,196 @@ +import type { ChatCompletionResult, ChatMessage, ChatOptions, LLMUsage } from "../llm/base-provider.js"; +import { BaseLLMProvider } from "../llm/base-provider.js"; +import { logger } from "../../lib/logger.js"; + +// Admission keys identify one physical provider endpoint. Text work keys on the configured +// connection id; image work keys on the resolved base URL plus the endpoint id where the +// backend needs one, because most image callers have no connection row in scope and +// unregistered foreground work would defeat the priority rule. The keyspaces do not overlap. +type ConnectionState = { + foregroundActive: number; + backgroundActive: boolean; + lastForegroundFinishedAt: number; +}; + +const states = new Map(); +export const BACKGROUND_CONNECTION_IDLE_MS = 30_000; + +export type ConnectionAttemptOutcome = "completed" | "failed"; +export type ConnectionAttemptFinalizer = (outcome: ConnectionAttemptOutcome) => void | Promise; +export type ConnectionAdmissionMode = + | { kind: "foreground" } + | { + kind: "background"; + beforeAttempt?: () => void | ConnectionAttemptFinalizer | Promise; + }; + +export class BackgroundConnectionBusyError extends Error { + constructor(readonly connectionId: string) { + super(`Connection ${connectionId} is not available for background generation.`); + this.name = "BackgroundConnectionBusyError"; + } +} + +export class ConnectionAttemptRejectedError extends Error { + constructor(readonly cause: unknown) { + super("Connection attempt was rejected before provider work started.", { cause }); + this.name = "ConnectionAttemptRejectedError"; + } +} + +export class ConnectionAttemptFinalizationError extends Error { + constructor(readonly cause: unknown) { + super("Connection attempt accounting failed after provider work finished.", { cause }); + this.name = "ConnectionAttemptFinalizationError"; + } +} + +export function isConnectionAdmissionFailure(error: unknown): boolean { + return ( + error instanceof BackgroundConnectionBusyError || + error instanceof ConnectionAttemptRejectedError || + error instanceof ConnectionAttemptFinalizationError + ); +} + +function stateFor(connectionId: string): ConnectionState { + const existing = states.get(connectionId); + if (existing) return existing; + const state = { foregroundActive: 0, backgroundActive: false, lastForegroundFinishedAt: 0 }; + states.set(connectionId, state); + return state; +} + +export function beginForegroundConnection(connectionId: string): () => void { + const state = stateFor(connectionId); + state.foregroundActive += 1; + let released = false; + return () => { + if (released) return; + released = true; + state.foregroundActive -= 1; + state.lastForegroundFinishedAt = Date.now(); + }; +} + +export function tryBackgroundConnection( + connectionId: string, + at: Date, +): { acquired: false } | { acquired: true; release: () => void } { + const state = stateFor(connectionId); + if ( + state.backgroundActive || + state.foregroundActive > 0 || + at.getTime() - state.lastForegroundFinishedAt < BACKGROUND_CONNECTION_IDLE_MS + ) { + return { acquired: false }; + } + state.backgroundActive = true; + return { + acquired: true, + release: () => { + state.backgroundActive = false; + }, + }; +} + +export function resetConnectionAdmissionForTests(): void { + states.clear(); +} + +async function beginConnectionAttempt( + connectionId: string, + mode: ConnectionAdmissionMode, +): Promise<{ release: () => void; finalize?: ConnectionAttemptFinalizer }> { + if (mode.kind === "foreground") return { release: beginForegroundConnection(connectionId) }; + + const admission = tryBackgroundConnection(connectionId, new Date()); + if (!admission.acquired) throw new BackgroundConnectionBusyError(connectionId); + try { + return { release: admission.release, finalize: (await mode.beforeAttempt?.()) || undefined }; + } catch (error) { + admission.release(); + throw new ConnectionAttemptRejectedError(error); + } +} + +/** + * Record the attempt outcome without letting accounting failures overwrite a provider + * error: the operation's own failure is what the caller (and the fallback chain) needs to + * see. A finalization failure only surfaces when the operation itself succeeded. + */ +async function finalizeConnectionAttempt( + attempt: { release: () => void; finalize?: ConnectionAttemptFinalizer }, + outcome: ConnectionAttemptOutcome, +): Promise { + try { + await attempt.finalize?.(outcome); + } catch (error) { + if (outcome === "completed") throw new ConnectionAttemptFinalizationError(error); + logger.error(error, "[connection-admission] Attempt accounting failed after a failed provider call"); + } finally { + attempt.release(); + } +} + +export async function withConnectionAdmission( + connectionId: string, + mode: ConnectionAdmissionMode, + operation: () => Promise, +): Promise { + const attempt = await beginConnectionAttempt(connectionId, mode); + let outcome: ConnectionAttemptOutcome = "failed"; + try { + const result = await operation(); + outcome = "completed"; + return result; + } finally { + await finalizeConnectionAttempt(attempt, outcome); + } +} + +export class ConnectionAdmissionProvider extends BaseLLMProvider { + constructor( + readonly provider: BaseLLMProvider, + private readonly connectionId: string, + private readonly mode: ConnectionAdmissionMode = { kind: "foreground" }, + ) { + super("", "", provider.maxContextValue ?? undefined, null, provider.maxTokensOverrideValue); + } + + async *chat(messages: ChatMessage[], options: ChatOptions): AsyncGenerator { + const attempt = await beginConnectionAttempt(this.connectionId, this.mode); + let outcome: ConnectionAttemptOutcome = "failed"; + try { + const result = yield* this.provider.chat(messages, options); + outcome = "completed"; + return result; + } finally { + await finalizeConnectionAttempt(attempt, outcome); + } + } + + async chatComplete(messages: ChatMessage[], options: ChatOptions): Promise { + return withConnectionAdmission(this.connectionId, this.mode, () => this.provider.chatComplete(messages, options)); + } + + async embed(texts: string[], model: string, signal?: AbortSignal): Promise { + return withConnectionAdmission(this.connectionId, this.mode, () => this.provider.embed(texts, model, signal)); + } +} + +/** + * Look through the admission decorator to the provider that actually talks to the backend. + * Diagnostics that report a provider class need the concrete one, not the wrapper. + */ +export function unwrapConnectionAdmissionProvider(provider: BaseLLMProvider): BaseLLMProvider { + return provider instanceof ConnectionAdmissionProvider ? provider.provider : provider; +} + +export function withConnectionAdmissionProvider( + provider: BaseLLMProvider, + connectionId: string, + mode: ConnectionAdmissionMode = { kind: "foreground" }, +): BaseLLMProvider { + return new ConnectionAdmissionProvider(provider, connectionId, mode); +} diff --git a/packages/server/src/services/image/image-generation.ts b/packages/server/src/services/image/image-generation.ts index 7cb5dc05fe..ddd4020537 100644 --- a/packages/server/src/services/image/image-generation.ts +++ b/packages/server/src/services/image/image-generation.ts @@ -30,6 +30,11 @@ import { generateRunPodComfyUI } from "./runpod-comfyui.service.js"; import { logger, logDebugOverride } from "../../lib/logger.js"; import { assertInsideDir, normalizeLoopbackUrl, safeFetch, validateOutboundUrl } from "../../utils/security.js"; import { notifyGenerationFallback, type GenerationFallbackNotifier } from "../generation/fallback-notification.js"; +import { + isConnectionAdmissionFailure, + withConnectionAdmission, + type ConnectionAdmissionMode, +} from "../generation/connection-admission.js"; import { COMFYUI_MAX_REFERENCE_IMAGES, findMissingComfyReferenceSlots, @@ -115,6 +120,8 @@ export interface ImageGenRequest { signal?: AbortSignal; /** Emit the final provider request even when the global log level is above debug. */ debugMode?: boolean; + /** Defaults to foreground: the caller is servicing a user-visible request. */ + admissionMode?: ConnectionAdmissionMode; /** Called immediately before a configured fallback connection is attempted. */ onFallback?: GenerationFallbackNotifier; /** Optional one-shot backup connection used only when the primary image request fails. */ @@ -194,6 +201,23 @@ function resolveImageBackend(source: string, baseUrl: string, serviceHint: strin const IMAGE_GEN_TIMEOUT = Number(process.env.IMAGE_GEN_TIMEOUT_MS ?? 1_800_000); const COMFYUI_GEN_TIMEOUT_SECONDS = Number(process.env.COMFYUI_GEN_TIMEOUT ?? 2400); +/** + * Identify the physical image endpoint an admission slot belongs to. RunPod connections share + * one API base URL and select the actual endpoint with a separate id, so the base URL alone + * would make two independent endpoints contend for a single slot. Every other backend's base + * URL already is the physical target, and a stale `imageEndpointId` left on an imported or + * copied connection must not split one ComfyUI/A1111 endpoint into separate slots. + */ +export function imageAdmissionKey( + normalizedBaseUrl: string, + resolvedSource: string, + imageEndpointId?: string, +): string { + if (resolvedSource !== "runpod_comfyui") return normalizedBaseUrl; + const endpointId = imageEndpointId?.trim(); + return endpointId ? `${normalizedBaseUrl}#${endpointId}` : normalizedBaseUrl; +} + /** * Generate an image using the configured image generation connection. * Returns the base64 data and metadata needed to save it. @@ -213,7 +237,7 @@ export async function generateImage( : IMAGE_GEN_TIMEOUT; try { - return await withImageGenerationDeadline(request, generationTimeoutMs, async (signal) => { + const physicalRequest = () => withImageGenerationDeadline(request, generationTimeoutMs, async (signal) => { const allowLocalUrls = request.allowLocalUrls ?? (await shouldAllowLocalUrlsForImageConnection(normalizedBaseUrl, resolvedSource)); const scopedRequest = { @@ -269,9 +293,20 @@ export async function generateImage( return generateOpenAI(normalizedBaseUrl, apiKey, scopedRequest); } }); + // Admit on the resolved endpoint rather than a connection id: every caller reaches this + // function, but only some have a connection row in scope, and foreground work that + // registers nothing would let background preparation start on top of it. + // ponytail: image work keys on the endpoint URL while text work keys on the connection + // id, so the two do not hold each other off on a connection used for both. Unify the + // key if that overlap ever shows up in practice. + return await withConnectionAdmission( + imageAdmissionKey(normalizedBaseUrl, resolvedSource, request.imageEndpointId), + request.admissionMode ?? { kind: "foreground" }, + physicalRequest, + ); } catch (error) { const fallback = request.fallback; - if (!fallback || request.signal?.aborted) throw error; + if (!fallback || request.signal?.aborted || isConnectionAdmissionFailure(error)) throw error; logger.warn( error, "[illustrator-fallback] Primary image generation failed; retrying with connection %s (%s)", diff --git a/packages/server/src/services/llm/connection-fallback-provider.ts b/packages/server/src/services/llm/connection-fallback-provider.ts index 1136c353b1..8a2d49ebcf 100644 --- a/packages/server/src/services/llm/connection-fallback-provider.ts +++ b/packages/server/src/services/llm/connection-fallback-provider.ts @@ -7,6 +7,11 @@ import { } from "../../routes/generate/generate-route-utils.js"; import { logger } from "../../lib/logger.js"; import { notifyGenerationFallback, type GenerationFallbackNotifier } from "../generation/fallback-notification.js"; +import { + isConnectionAdmissionFailure, + withConnectionAdmissionProvider, + type ConnectionAdmissionMode, +} from "../generation/connection-admission.js"; export type FallbackConnection = { id: string; @@ -34,6 +39,7 @@ type ConnectionFallbackProviderArgs = { fallbackBaseUrl: string; category: "main" | "agents"; onFallback?: GenerationFallbackNotifier; + admissionMode?: ConnectionAdmissionMode; }; function isEnabled(value: unknown): boolean { @@ -148,16 +154,28 @@ export class ConnectionFallbackProvider extends BaseLLMProvider { } : options; const generation = this.primary.chat(messages, primaryOptions); - let result = await generation.next(); - while (!result.done) { - emittedUsableOutput ||= result.value.trim().length > 0; - yield result.value; - result = await generation.next(); + try { + let result = await generation.next(); + while (!result.done) { + emittedUsableOutput ||= result.value.trim().length > 0; + yield result.value; + result = await generation.next(); + } + if (emittedUsableOutput || options.signal?.aborted) return result.value; + } finally { + // Drive the primary to completion if our consumer abandoned us mid-stream. The manual + // loop above does not forward an early return the way `yield*` would, so without this + // an admission wrapper around the primary never runs its own finally and leaks the + // connection's foreground slot for the lifetime of the process. A cleanup rejection is + // logged rather than thrown: the provider error the catch below inspects is what decides + // whether we fall back, and it must not be replaced by a teardown failure. + await generation.return(undefined).catch((closeError: unknown) => { + logger.warn(closeError, "[%s-fallback] Failed to close the primary generation stream", this.category); + }); } - if (emittedUsableOutput || options.signal?.aborted) return result.value; await this.logFallback(new Error("Primary provider returned an empty completion")); } catch (error) { - if (emittedUsableOutput || isAbortFailure(error, options.signal)) throw error; + if (emittedUsableOutput || isAbortFailure(error, options.signal) || isConnectionAdmissionFailure(error)) throw error; await this.logFallback(error); } options.signal?.throwIfAborted(); @@ -171,7 +189,7 @@ export class ConnectionFallbackProvider extends BaseLLMProvider { if (hasUsableOutput || options.signal?.aborted) return result; await this.logFallback(new Error("Primary provider returned an empty completion")); } catch (error) { - if (isAbortFailure(error, options.signal)) throw error; + if (isAbortFailure(error, options.signal) || isConnectionAdmissionFailure(error)) throw error; await this.logFallback(error); } options.signal?.throwIfAborted(); @@ -190,26 +208,31 @@ export function withConnectionFallbackProvider({ fallbackBaseUrl, category, onFallback, + admissionMode = { kind: "foreground" }, }: ConnectionFallbackProviderArgs): BaseLLMProvider { + const admittedPrimary = withConnectionAdmissionProvider(primary, primaryConnectionId, admissionMode); if ( !fallbackConnection || fallbackConnection.id === primaryConnectionId || !fallbackConnection.model?.trim() || !fallbackBaseUrl ) { - return primary; + return admittedPrimary; } - - const fallback = createLLMProvider( - fallbackConnection.provider, - fallbackBaseUrl, - fallbackConnection.apiKey, - fallbackConnection.maxContext, - fallbackConnection.openrouterProvider, - fallbackConnection.maxTokensOverride, - isEnabled(fallbackConnection.claudeFastMode), - isEnabled(fallbackConnection.treatAsLocalEndpoint), - fallbackConnection.defaultParameters, + const fallback = withConnectionAdmissionProvider( + createLLMProvider( + fallbackConnection.provider, + fallbackBaseUrl, + fallbackConnection.apiKey, + fallbackConnection.maxContext, + fallbackConnection.openrouterProvider, + fallbackConnection.maxTokensOverride, + isEnabled(fallbackConnection.claudeFastMode), + isEnabled(fallbackConnection.treatAsLocalEndpoint), + fallbackConnection.defaultParameters, + ), + fallbackConnection.id, + admissionMode, ); - return new ConnectionFallbackProvider(primary, fallback, fallbackConnection, category, onFallback); + return new ConnectionFallbackProvider(admittedPrimary, fallback, fallbackConnection, category, onFallback); } diff --git a/packages/server/src/services/llm/provider-registry.ts b/packages/server/src/services/llm/provider-registry.ts index aea325600b..2cb97a8767 100644 --- a/packages/server/src/services/llm/provider-registry.ts +++ b/packages/server/src/services/llm/provider-registry.ts @@ -9,6 +9,7 @@ import { GrokSubscriptionProvider } from "./providers/grok-subscription.provider import { GoogleProvider } from "./providers/google.provider.js"; import type { BaseLLMProvider } from "./base-provider.js"; import { withConnectionDefaultParameters } from "./connection-default-provider.js"; +import { withConnectionAdmissionProvider } from "../generation/connection-admission.js"; export function normalizeCohereOpenAIBaseUrl(baseUrl: string): string { const trimmed = baseUrl.replace(/\/+$/, ""); @@ -48,6 +49,8 @@ export function createLLMProvider( treatAsLocalEndpoint?: boolean, /** Stored connection defaults. Custom Parameters are bound to every text request made by this provider. */ defaultParameters?: unknown, + /** Configured connection ID for direct foreground calls. Fallback wrappers admit their providers separately. */ + connectionId?: string, ): BaseLLMProvider { const normalizedMaxContext = typeof maxContext === "number" && Number.isFinite(maxContext) && maxContext > 0 @@ -159,5 +162,6 @@ export function createLLMProvider( ); break; } - return withConnectionDefaultParameters(resolved, defaultParameters); + const configured = withConnectionDefaultParameters(resolved, defaultParameters); + return connectionId ? withConnectionAdmissionProvider(configured, connectionId) : configured; } diff --git a/packages/server/src/services/memory-recall-embedding.ts b/packages/server/src/services/memory-recall-embedding.ts index a9c78f7c08..777218d22a 100644 --- a/packages/server/src/services/memory-recall-embedding.ts +++ b/packages/server/src/services/memory-recall-embedding.ts @@ -164,6 +164,7 @@ export async function resolveMemoryRecallEmbeddingSource( embeddingConnection.claudeFastMode === "true", embeddingConnection.treatAsLocalEndpoint === "true", embeddingConnection.defaultParameters, + embeddingConnection.id, ); const label = `${embeddingConnection.name || embeddingConnection.provider} (${embeddingModel})`; diff --git a/packages/server/src/services/professor-mari/workspace-agent.service.ts b/packages/server/src/services/professor-mari/workspace-agent.service.ts index 02d0a43bb6..cf041624c3 100644 --- a/packages/server/src/services/professor-mari/workspace-agent.service.ts +++ b/packages/server/src/services/professor-mari/workspace-agent.service.ts @@ -840,6 +840,7 @@ function createProviderForConnection(connection: WorkspaceConnection): BaseLLMPr bool(connection.claudeFastMode), bool(connection.treatAsLocalEndpoint), connection.defaultParameters, + connection.id, ); } diff --git a/packages/server/src/services/spotify/dj-mari-playlist.service.ts b/packages/server/src/services/spotify/dj-mari-playlist.service.ts index 7f7896ae10..8870e4a3a0 100644 --- a/packages/server/src/services/spotify/dj-mari-playlist.service.ts +++ b/packages/server/src/services/spotify/dj-mari-playlist.service.ts @@ -593,6 +593,7 @@ async function generatePlaylistPlan(args: { conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ); const userContext = { diff --git a/packages/server/src/services/turn-games/turn-game-bot-runner.service.ts b/packages/server/src/services/turn-games/turn-game-bot-runner.service.ts index 5e17c40d01..a55a724cff 100644 --- a/packages/server/src/services/turn-games/turn-game-bot-runner.service.ts +++ b/packages/server/src/services/turn-games/turn-game-bot-runner.service.ts @@ -274,6 +274,7 @@ export async function runTurnGameBotTurns(args: RunBotTurnsArgs): Promise conn.claudeFastMode === "true", conn.treatAsLocalEndpoint === "true", conn.defaultParameters, + conn.id, ) : null); if (!provider) { diff --git a/scripts/regressions/provider-compat.regression.ts b/scripts/regressions/provider-compat.regression.ts index baa028f48e..c3ce2abda1 100644 --- a/scripts/regressions/provider-compat.regression.ts +++ b/scripts/regressions/provider-compat.regression.ts @@ -54,6 +54,14 @@ import { } from "../../packages/server/src/services/generation/fallback-notification.js"; import { resolveStoredChatOptions } from "../../packages/server/src/services/generation/generation-parameters.js"; import { resolveMainGenerationToolChoice } from "../../packages/server/src/services/generation/tool-resolution-runtime.js"; +import { imageAdmissionKey } from "../../packages/server/src/services/image/image-generation.js"; +import { + BACKGROUND_CONNECTION_IDLE_MS, + ConnectionAttemptRejectedError, + resetConnectionAdmissionForTests, + tryBackgroundConnection, + withConnectionAdmissionProvider, +} from "../../packages/server/src/services/generation/connection-admission.js"; class RegressionProvider extends BaseLLMProvider { calls = 0; @@ -759,6 +767,151 @@ const fallbackConnection: FallbackConnection = { }, }), }; + +resetConnectionAdmissionForTests(); +let releasePrimaryCall!: () => void; +const primaryCallHeld = new Promise((resolve) => { + releasePrimaryCall = resolve; +}); +class HeldProvider extends BaseLLMProvider { + started!: () => void; + readonly startedPromise = new Promise((resolve) => { + this.started = resolve; + }); + + constructor(private readonly held: Promise, private readonly failure?: Error) { + super("", ""); + } + + async *chat(): AsyncGenerator { + this.started(); + await this.held; + if (this.failure) throw this.failure; + yield "held response"; + } +} +const heldPrimary = new HeldProvider(primaryCallHeld, new Error("primary unavailable")); +const admittedPrimary = withConnectionAdmissionProvider(heldPrimary, "primary-connection"); +const primaryRun = collectProviderOutput(admittedPrimary, { model: "primary-model" }); +await heldPrimary.startedPromise; +assert.equal(tryBackgroundConnection("primary-connection", new Date()).acquired, false); +const unrelatedFallbackAdmission = tryBackgroundConnection("fallback-connection", new Date()); +assert.equal(unrelatedFallbackAdmission.acquired, true, "primary work must not occupy the fallback connection"); +if (unrelatedFallbackAdmission.acquired) unrelatedFallbackAdmission.release(); +releasePrimaryCall(); +await assert.rejects(primaryRun, /primary unavailable/); +const releasedPrimaryAdmission = tryBackgroundConnection( + "primary-connection", + new Date(Date.now() + BACKGROUND_CONNECTION_IDLE_MS), +); +assert.equal(releasedPrimaryAdmission.acquired, true, "a failed primary call must release its foreground slot"); +if (releasedPrimaryAdmission.acquired) releasedPrimaryAdmission.release(); + +let releaseFallbackCall!: () => void; +const fallbackCallHeld = new Promise((resolve) => { + releaseFallbackCall = resolve; +}); +const heldFallback = new HeldProvider(fallbackCallHeld); +const admittedFallback = withConnectionAdmissionProvider(heldFallback, "fallback-connection"); +const fallbackRun = collectProviderOutput(admittedFallback, { model: "fallback-model" }); +await heldFallback.startedPromise; +assert.equal(tryBackgroundConnection("fallback-connection", new Date()).acquired, false); +const unrelatedPrimaryAdmission = tryBackgroundConnection("other-primary-connection", new Date()); +assert.equal(unrelatedPrimaryAdmission.acquired, true, "fallback work must not occupy a primary connection"); +if (unrelatedPrimaryAdmission.acquired) unrelatedPrimaryAdmission.release(); +releaseFallbackCall(); +assert.equal(await fallbackRun, "held response"); +const releasedFallbackAdmission = tryBackgroundConnection( + "fallback-connection", + new Date(Date.now() + BACKGROUND_CONNECTION_IDLE_MS), +); +assert.equal(releasedFallbackAdmission.acquired, true, "a completed call must release its foreground slot"); +if (releasedFallbackAdmission.acquired) releasedFallbackAdmission.release(); + +// A consumer that stops reading mid-stream must not strand the primary connection's +// foreground slot: the fallback provider drains the primary manually, so an early return +// has to be forwarded explicitly or background work is locked out until a restart. +resetConnectionAdmissionForTests(); +class EndlessProvider extends BaseLLMProvider { + constructor() { + super("", ""); + } + + async *chat(): AsyncGenerator { + while (true) yield "chunk"; + } +} +const abandonedProvider = new ConnectionFallbackProvider( + withConnectionAdmissionProvider(new EndlessProvider(), "abandoned-connection"), + new RegressionProvider(["unused fallback"]), + fallbackConnection, + "main", +); +const abandonedStream = abandonedProvider.chat([{ role: "user", content: "hi" }], { model: "abandoned-model" }); +assert.equal((await abandonedStream.next()).value, "chunk"); +await abandonedStream.return(undefined); +assert.equal( + tryBackgroundConnection("abandoned-connection", new Date(Date.now() + BACKGROUND_CONNECTION_IDLE_MS)).acquired, + true, + "abandoning the stream must release the primary connection's foreground slot", +); + +// RunPod connections share one API base URL and pick the physical endpoint with a separate +// id, so the key has to carry that id or two independent endpoints contend for one slot. +assert.notEqual( + imageAdmissionKey("https://api.runpod.ai/v2", "runpod_comfyui", "abc123"), + imageAdmissionKey("https://api.runpod.ai/v2", "runpod_comfyui", "def456"), + "RunPod endpoints sharing a base URL must not share an admission key", +); +assert.equal( + imageAdmissionKey("https://api.runpod.ai/v2", "runpod_comfyui", "abc123"), + imageAdmissionKey("https://api.runpod.ai/v2", "runpod_comfyui", " abc123 "), + "endpoint ids must be compared after trimming", +); +assert.equal( + imageAdmissionKey("http://127.0.0.1:8188", "comfyui"), + "http://127.0.0.1:8188", + "backends whose base URL is the physical target keep the bare URL as their key", +); +// A stale endpoint id on an imported/copied non-RunPod connection must not split one +// physical endpoint into two admission slots. +assert.equal( + imageAdmissionKey("http://127.0.0.1:8188", "comfyui", "abc123"), + imageAdmissionKey("http://127.0.0.1:8188", "comfyui"), + "non-RunPod backends must ignore a leftover endpoint id", +); +assert.equal( + imageAdmissionKey("http://127.0.0.1:7860", "automatic1111", "def456"), + "http://127.0.0.1:7860", + "A1111 admission keys stay bare regardless of endpoint id", +); + +let backgroundAttempts = 0; +const backgroundOutcomes: string[] = []; +const admittedRepeated = withConnectionAdmissionProvider(new RegressionProvider(["ok"]), "background-connection", { + kind: "background", + beforeAttempt: () => { + backgroundAttempts += 1; + return (outcome) => { + backgroundOutcomes.push(outcome); + }; + }, +}); +await admittedRepeated.chatComplete([{ role: "user", content: "first" }], { model: "model" }); +await admittedRepeated.chatComplete([{ role: "user", content: "correction" }], { model: "model" }); +assert.equal(backgroundAttempts, 2, "each physical correction/retry call must claim separately"); +assert.deepEqual(backgroundOutcomes, ["completed", "completed"]); + +const rejectedAttempt = withConnectionAdmissionProvider(new RegressionProvider(["unused"]), "rejected-connection", { + kind: "background", + beforeAttempt: () => { + throw new Error("attempt budget exhausted"); + }, +}); +await assert.rejects( + rejectedAttempt.chatComplete([{ role: "user", content: "test" }], { model: "model" }), + (error) => error instanceof ConnectionAttemptRejectedError && error.cause instanceof Error && /budget exhausted/.test(error.cause.message), +); const primaryFailure = new RegressionProvider([], new Error("primary unavailable")); const successfulFallback = new RegressionProvider(["fallback response"]); const fallbackProvider = new ConnectionFallbackProvider(primaryFailure, successfulFallback, fallbackConnection, "main");