Spaces:

automindy
/

fokemal

Sleeping

fokemal / server /src /__tests__ /lib /fallback-loop-lease.test.ts

Sync to upstream v0.13.7, keep HF Space adaptations

dcedb25 22 days ago

4.44 kB

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { initDb } from '../../db/index.js';
import { runFallbackLoop, newFallbackState } from '../../lib/fallback-loop.js';
import { acquireLease, releaseLease, resetLeases, inFlightForKey } from '../../services/ratelimit.js';
import type { RouteResult } from '../../services/router.js';

/**
 * The loop must free a route's in-flight lease on every exit path. A leak would
 * leave the key's concurrency budget permanently short (until the age-out
 * backstop), so this covers each way an iteration can end.
 */
function leasedRoute(keyId: number): RouteResult {
  const id = acquireLease('groq', 'llama-3.3-70b', keyId, 100);
  return {
    provider: {} as any,
    modelId: 'llama-3.3-70b',
    modelDbId: 1,
    apiKey: 'test-key',
    keyId,
    platform: 'groq',
    displayName: 'Llama 3.3 70B',
    rpdLimit: null,
    tpdLimit: null,
    release: () => releaseLease(id),
  };
}

function hooks(overrides: Record<string, unknown> = {}) {
  return {
    state: newFallbackState(),
    attemptLog: [],
    route: () => leasedRoute(1),
    dispatch: async () => 'done' as const,
    logFailure: () => undefined,
    onFatal: () => undefined,
    onRoutingExhausted: () => undefined,
    onExhausted: () => undefined,
    ...overrides,
  } as any;
}

describe('fallback loop releases in-flight leases', () => {
  beforeEach(() => {
    initDb(':memory:');
    resetLeases();
    vi.spyOn(console, 'log').mockImplementation(() => {});
    vi.spyOn(console, 'error').mockImplementation(() => {});
    vi.spyOn(console, 'warn').mockImplementation(() => {});
  });

  afterEach(() => {
    resetLeases();
    vi.restoreAllMocks();
  });

  it('releases after a successful dispatch', async () => {
    await runFallbackLoop(hooks());
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('releases after a committed stream', async () => {
    await runFallbackLoop(hooks({ dispatch: async () => 'committed' as const }));
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('releases every attempt when retryable failures exhaust the budget', async () => {
    let dispatched = 0;
    await runFallbackLoop(hooks({
      maxRetries: 4,
      dispatch: async () => {
        dispatched++;
        throw Object.assign(new Error('Groq API error 429: rate limit'), { status: 429 });
      },
    }));

    expect(dispatched).toBe(4);
    // Four leases were taken across the four attempts; none may survive.
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('releases on an auth failure that rotates the key', async () => {
    await runFallbackLoop(hooks({
      maxRetries: 2,
      dispatch: async () => {
        throw Object.assign(new Error('Groq API error 401: Invalid API Key'), { status: 401 });
      },
    }));
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('releases on a non-retryable fatal error', async () => {
    const onFatal = vi.fn();
    await runFallbackLoop(hooks({
      onFatal,
      dispatch: async () => { throw Object.assign(new Error('boom'), { status: 418 }); },
    }));
    expect(onFatal).toHaveBeenCalled();
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('releases when dispatch violates the outcome contract', async () => {
    const onFatal = vi.fn();
    await runFallbackLoop(hooks({ onFatal, dispatch: async () => 'nonsense' as any }));
    expect(onFatal).toHaveBeenCalled();
    expect(inFlightForKey('groq', 1)).toBe(0);
  });

  it('does not mask the provider error if release itself throws', async () => {
    const onFatal = vi.fn();
    const bad = {
      ...leasedRoute(1),
      release: () => { throw new Error('release exploded'); },
    };
    await runFallbackLoop(hooks({
      route: () => bad,
      onFatal,
      dispatch: async () => { throw Object.assign(new Error('real provider failure'), { status: 418 }); },
    })).catch(() => undefined);

    // onFatal must have seen the provider error, not the release error.
    expect(onFatal).toHaveBeenCalled();
    const err = onFatal.mock.calls[0]![1] as Error;
    expect(err.message).toContain('real provider failure');
  });

  it('tolerates a route with no release function', async () => {
    const noRelease = { ...leasedRoute(2) } as RouteResult;
    delete (noRelease as { release?: unknown }).release;
    await expect(runFallbackLoop(hooks({ route: () => noRelease }))).resolves.toBeUndefined();
  });
});

	import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
	import { initDb } from '../../db/index.js';
	import { runFallbackLoop, newFallbackState } from '../../lib/fallback-loop.js';
	import { acquireLease, releaseLease, resetLeases, inFlightForKey } from '../../services/ratelimit.js';
	import type { RouteResult } from '../../services/router.js';

	/**
	* The loop must free a route's in-flight lease on every exit path. A leak would
	* leave the key's concurrency budget permanently short (until the age-out
	* backstop), so this covers each way an iteration can end.
	*/
	function leasedRoute(keyId: number): RouteResult {
	const id = acquireLease('groq', 'llama-3.3-70b', keyId, 100);
	return {
	provider: {} as any,
	modelId: 'llama-3.3-70b',
	modelDbId: 1,
	apiKey: 'test-key',
	keyId,
	platform: 'groq',
	displayName: 'Llama 3.3 70B',
	rpdLimit: null,
	tpdLimit: null,
	release: () => releaseLease(id),
	};
	}

	function hooks(overrides: Record<string, unknown> = {}) {
	return {
	state: newFallbackState(),
	attemptLog: [],
	route: () => leasedRoute(1),
	dispatch: async () => 'done' as const,
	logFailure: () => undefined,
	onFatal: () => undefined,
	onRoutingExhausted: () => undefined,
	onExhausted: () => undefined,
	...overrides,
	} as any;
	}

	describe('fallback loop releases in-flight leases', () => {
	beforeEach(() => {
	initDb(':memory:');
	resetLeases();
	vi.spyOn(console, 'log').mockImplementation(() => {});
	vi.spyOn(console, 'error').mockImplementation(() => {});
	vi.spyOn(console, 'warn').mockImplementation(() => {});
	});

	afterEach(() => {
	resetLeases();
	vi.restoreAllMocks();
	});

	it('releases after a successful dispatch', async () => {
	await runFallbackLoop(hooks());
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('releases after a committed stream', async () => {
	await runFallbackLoop(hooks({ dispatch: async () => 'committed' as const }));
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('releases every attempt when retryable failures exhaust the budget', async () => {
	let dispatched = 0;
	await runFallbackLoop(hooks({
	maxRetries: 4,
	dispatch: async () => {
	dispatched++;
	throw Object.assign(new Error('Groq API error 429: rate limit'), { status: 429 });
	},
	}));

	expect(dispatched).toBe(4);
	// Four leases were taken across the four attempts; none may survive.
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('releases on an auth failure that rotates the key', async () => {
	await runFallbackLoop(hooks({
	maxRetries: 2,
	dispatch: async () => {
	throw Object.assign(new Error('Groq API error 401: Invalid API Key'), { status: 401 });
	},
	}));
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('releases on a non-retryable fatal error', async () => {
	const onFatal = vi.fn();
	await runFallbackLoop(hooks({
	onFatal,
	dispatch: async () => { throw Object.assign(new Error('boom'), { status: 418 }); },
	}));
	expect(onFatal).toHaveBeenCalled();
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('releases when dispatch violates the outcome contract', async () => {
	const onFatal = vi.fn();
	await runFallbackLoop(hooks({ onFatal, dispatch: async () => 'nonsense' as any }));
	expect(onFatal).toHaveBeenCalled();
	expect(inFlightForKey('groq', 1)).toBe(0);
	});

	it('does not mask the provider error if release itself throws', async () => {
	const onFatal = vi.fn();
	const bad = {
	...leasedRoute(1),
	release: () => { throw new Error('release exploded'); },
	};
	await runFallbackLoop(hooks({
	route: () => bad,
	onFatal,
	dispatch: async () => { throw Object.assign(new Error('real provider failure'), { status: 418 }); },
	})).catch(() => undefined);

	// onFatal must have seen the provider error, not the release error.
	expect(onFatal).toHaveBeenCalled();
	const err = onFatal.mock.calls[0]![1] as Error;
	expect(err.message).toContain('real provider failure');
	});

	it('tolerates a route with no release function', async () => {
	const noRelease = { ...leasedRoute(2) } as RouteResult;
	delete (noRelease as { release?: unknown }).release;
	await expect(runFallbackLoop(hooks({ route: () => noRelease }))).resolves.toBeUndefined();
	});
	});