mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-05-13 07:46:47 +00:00
🔑 feat: Per-User Google SA For Claude On Vertex
When `GOOGLE_KEY=user_provided` and the user has registered a Google service-account JSON via the existing Gemini per-user-keys path, reuse that same SA for Anthropic-Vertex (`ANTHROPIC_USE_VERTEX=true`) instead of forcing every user onto the global `GOOGLE_SERVICE_KEY_FILE`. Falls back to the file/env-based loader on no-stored-key, parse error, or malformed result so the existing global Vertex path keeps working for users who have not registered a key. Closes #12979 Signed-off-by: ChrisJr404 <chris@hacknow.com>
This commit is contained in:
ChrisJr404
parent
6b5596ec36
commit
d28f521cb5
2 changed files with 345 additions and 6 deletions
274
packages/api/src/endpoints/anthropic/initialize.spec.ts
Normal file
274
packages/api/src/endpoints/anthropic/initialize.spec.ts
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
import { AuthKeys, EModelEndpoint, ErrorTypes } from 'librechat-data-provider';
|
||||
import type { BaseInitializeParams } from '~/types';
|
||||
|
||||
const mockLoadAnthropicVertexCredentials = jest.fn();
|
||||
const mockGetVertexCredentialOptions = jest.fn();
|
||||
jest.mock('./vertex', () => ({
|
||||
loadAnthropicVertexCredentials: (...args: unknown[]) =>
|
||||
mockLoadAnthropicVertexCredentials(...args),
|
||||
getVertexCredentialOptions: (...args: unknown[]) => mockGetVertexCredentialOptions(...args),
|
||||
}));
|
||||
|
||||
const mockGetLLMConfig = jest
|
||||
.fn()
|
||||
.mockReturnValue({ llmConfig: { model: 'claude-3-7-sonnet-20250219' } });
|
||||
jest.mock('./llm', () => ({
|
||||
getLLMConfig: (...args: unknown[]) => mockGetLLMConfig(...args),
|
||||
}));
|
||||
|
||||
const mockLoadServiceKey = jest.fn();
|
||||
jest.mock('~/utils', () => ({
|
||||
isEnabled: (val?: unknown) => val === 'true' || val === true,
|
||||
loadServiceKey: (...args: unknown[]) => mockLoadServiceKey(...args),
|
||||
checkUserKeyExpiry: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
import { initializeAnthropic } from './initialize';
|
||||
|
||||
const GLOBAL_SERVICE_KEY = {
|
||||
project_id: 'global-project',
|
||||
client_email: 'global@global-project.iam.gserviceaccount.com',
|
||||
private_key: 'global-private-key',
|
||||
};
|
||||
|
||||
const USER_SERVICE_KEY = {
|
||||
project_id: 'user-project',
|
||||
client_email: 'user@user-project.iam.gserviceaccount.com',
|
||||
private_key: 'user-private-key',
|
||||
};
|
||||
|
||||
function createParams(
|
||||
env: Record<string, string | undefined>,
|
||||
dbOverrides: Partial<BaseInitializeParams['db']> = {},
|
||||
userId: string | null = 'user-1',
|
||||
): BaseInitializeParams & { _restore: () => void } {
|
||||
const savedEnv: Record<string, string | undefined> = {};
|
||||
for (const key of Object.keys(env)) {
|
||||
savedEnv[key] = process.env[key];
|
||||
if (env[key] === undefined) {
|
||||
delete process.env[key];
|
||||
} else {
|
||||
process.env[key] = env[key];
|
||||
}
|
||||
}
|
||||
|
||||
const db = {
|
||||
getUserKey: jest.fn(),
|
||||
getUserKeyValues: jest.fn(),
|
||||
...dbOverrides,
|
||||
} as unknown as BaseInitializeParams['db'];
|
||||
|
||||
const params: BaseInitializeParams = {
|
||||
req: {
|
||||
user: userId ? { id: userId } : undefined,
|
||||
body: {},
|
||||
config: { endpoints: {} },
|
||||
} as unknown as BaseInitializeParams['req'],
|
||||
endpoint: EModelEndpoint.anthropic,
|
||||
model_parameters: { model: 'claude-3-7-sonnet-20250219' },
|
||||
db,
|
||||
};
|
||||
|
||||
const restore = () => {
|
||||
for (const key of Object.keys(env)) {
|
||||
if (savedEnv[key] === undefined) {
|
||||
delete process.env[key];
|
||||
} else {
|
||||
process.env[key] = savedEnv[key];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
return Object.assign(params, { _restore: restore });
|
||||
}
|
||||
|
||||
describe('initializeAnthropic – per-user Vertex service-account credentials', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockLoadAnthropicVertexCredentials.mockResolvedValue({
|
||||
[AuthKeys.GOOGLE_SERVICE_KEY]: GLOBAL_SERVICE_KEY,
|
||||
});
|
||||
mockGetVertexCredentialOptions.mockReturnValue({});
|
||||
});
|
||||
|
||||
it('uses the user-stored Google SA when GOOGLE_KEY=user_provided and a valid key is registered', async () => {
|
||||
const storedUserKey = JSON.stringify({
|
||||
[AuthKeys.GOOGLE_SERVICE_KEY]: 'stringified-sa-json-from-storage',
|
||||
});
|
||||
mockLoadServiceKey.mockResolvedValueOnce(USER_SERVICE_KEY);
|
||||
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn().mockResolvedValue(storedUserKey),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(params.db.getUserKey).toHaveBeenCalledWith({
|
||||
userId: 'user-1',
|
||||
name: EModelEndpoint.google,
|
||||
});
|
||||
expect(mockLoadServiceKey).toHaveBeenCalledWith('stringified-sa-json-from-storage');
|
||||
expect(mockLoadAnthropicVertexCredentials).not.toHaveBeenCalled();
|
||||
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(USER_SERVICE_KEY);
|
||||
});
|
||||
|
||||
it('falls back to global file-based service key when the user has no stored Google key', async () => {
|
||||
const noUserKeyErr = new Error(JSON.stringify({ type: ErrorTypes.NO_USER_KEY }));
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn().mockRejectedValue(noUserKeyErr),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(params.db.getUserKey).toHaveBeenCalled();
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(GLOBAL_SERVICE_KEY);
|
||||
});
|
||||
|
||||
it('falls back to global service key when the stored Google entry is not valid JSON', async () => {
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn().mockResolvedValue('not-json'),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(mockLoadServiceKey).not.toHaveBeenCalled();
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(GLOBAL_SERVICE_KEY);
|
||||
});
|
||||
|
||||
it('falls back to global service key when loadServiceKey returns a malformed result', async () => {
|
||||
const storedUserKey = JSON.stringify({
|
||||
[AuthKeys.GOOGLE_SERVICE_KEY]: 'stringified-sa-json',
|
||||
});
|
||||
mockLoadServiceKey.mockResolvedValueOnce({ project_id: 'only-project-no-private-key' });
|
||||
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn().mockResolvedValue(storedUserKey),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(GLOBAL_SERVICE_KEY);
|
||||
});
|
||||
|
||||
it('does not attempt per-user lookup when GOOGLE_KEY is not set to user_provided', async () => {
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'AIza-some-static-google-api-key',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn(),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(params.db.getUserKey).not.toHaveBeenCalled();
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(GLOBAL_SERVICE_KEY);
|
||||
});
|
||||
|
||||
it('does not attempt per-user lookup when there is no authenticated user', async () => {
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{ getUserKey: jest.fn() },
|
||||
null,
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(params.db.getUserKey).not.toHaveBeenCalled();
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('handles non-NO_USER_KEY db errors by falling back gracefully', async () => {
|
||||
const params = createParams(
|
||||
{
|
||||
ANTHROPIC_USE_VERTEX: 'true',
|
||||
GOOGLE_KEY: 'user_provided',
|
||||
},
|
||||
{
|
||||
getUserKey: jest.fn().mockRejectedValue(new Error('mongo down')),
|
||||
},
|
||||
);
|
||||
|
||||
try {
|
||||
await initializeAnthropic(params);
|
||||
} finally {
|
||||
params._restore();
|
||||
}
|
||||
|
||||
expect(mockLoadAnthropicVertexCredentials).toHaveBeenCalledTimes(1);
|
||||
const [credentials] = mockGetLLMConfig.mock.calls[0] as [Record<string, unknown>];
|
||||
expect(credentials[AuthKeys.GOOGLE_SERVICE_KEY]).toEqual(GLOBAL_SERVICE_KEY);
|
||||
});
|
||||
});
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
import { EModelEndpoint, AuthKeys } from 'librechat-data-provider';
|
||||
import { EModelEndpoint, AuthKeys, ErrorTypes } from 'librechat-data-provider';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import type { BaseInitializeParams, InitializeResultBase, AnthropicConfigOptions } from '~/types';
|
||||
import { checkUserKeyExpiry, isEnabled } from '~/utils';
|
||||
import { checkUserKeyExpiry, isEnabled, loadServiceKey } from '~/utils';
|
||||
import { loadAnthropicVertexCredentials, getVertexCredentialOptions } from './vertex';
|
||||
import { getLLMConfig } from './llm';
|
||||
|
||||
|
|
@ -20,7 +21,7 @@ export async function initializeAnthropic({
|
|||
}: BaseInitializeParams): Promise<InitializeResultBase> {
|
||||
void endpoint;
|
||||
const appConfig = req.config;
|
||||
const { ANTHROPIC_API_KEY, ANTHROPIC_REVERSE_PROXY, PROXY } = process.env;
|
||||
const { ANTHROPIC_API_KEY, ANTHROPIC_REVERSE_PROXY, PROXY, GOOGLE_KEY } = process.env;
|
||||
const { key: expiresAt } = req.body;
|
||||
|
||||
let credentials: Record<string, unknown> = {};
|
||||
|
|
@ -35,9 +36,73 @@ export async function initializeAnthropic({
|
|||
(vertexConfig && vertexConfig.enabled !== false) || isEnabled(process.env.ANTHROPIC_USE_VERTEX);
|
||||
|
||||
if (useVertexAI) {
|
||||
// Load credentials with optional YAML config overrides
|
||||
const credentialOptions = vertexConfig ? getVertexCredentialOptions(vertexConfig) : undefined;
|
||||
credentials = await loadAnthropicVertexCredentials(credentialOptions);
|
||||
/**
|
||||
* When Gemini is configured for per-user service-account keys (GOOGLE_KEY=user_provided),
|
||||
* reuse the user's stored Google SA JSON for Claude-via-Vertex too. This matches the
|
||||
* pattern Gemini already follows in `initializeGoogle` and lets a single per-user SA
|
||||
* back both Gemini and Claude on Vertex.
|
||||
*
|
||||
* Falls back to the file/env-based loader on miss, parse error, or any other failure
|
||||
* so the global config path keeps working when a user has not registered a key.
|
||||
*/
|
||||
const userId = req.user?.id;
|
||||
const allowPerUserVertex = GOOGLE_KEY === 'user_provided' && !!userId;
|
||||
let userServiceKeyLoaded = false;
|
||||
|
||||
if (allowPerUserVertex) {
|
||||
try {
|
||||
const stored = await db.getUserKey({
|
||||
userId: userId as string,
|
||||
name: EModelEndpoint.google,
|
||||
});
|
||||
|
||||
let parsed: Record<string, unknown> | null = null;
|
||||
try {
|
||||
parsed = JSON.parse(stored) as Record<string, unknown>;
|
||||
} catch (parseErr) {
|
||||
logger.warn(
|
||||
'[initializeAnthropic] Failed to parse stored Google key for per-user Vertex; falling back to global service key',
|
||||
parseErr,
|
||||
);
|
||||
}
|
||||
|
||||
const sa = parsed?.[AuthKeys.GOOGLE_SERVICE_KEY];
|
||||
if (typeof sa === 'string' && sa.trim() !== '') {
|
||||
const loaded = await loadServiceKey(sa);
|
||||
if (loaded?.private_key && loaded?.project_id) {
|
||||
credentials[AuthKeys.GOOGLE_SERVICE_KEY] = loaded;
|
||||
userServiceKeyLoaded = true;
|
||||
} else {
|
||||
logger.warn(
|
||||
'[initializeAnthropic] Stored Google SA key for user missing private_key or project_id; falling back to global service key',
|
||||
);
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
// NO_USER_KEY: user has not registered a per-user SA; fall back silently.
|
||||
// Other errors: log but continue with the global service key so we degrade gracefully.
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
let isNoUserKey = false;
|
||||
try {
|
||||
const parsed = JSON.parse(message) as { type?: string };
|
||||
isNoUserKey = parsed?.type === ErrorTypes.NO_USER_KEY;
|
||||
} catch {
|
||||
// not a structured error, ignore
|
||||
}
|
||||
if (!isNoUserKey) {
|
||||
logger.warn(
|
||||
'[initializeAnthropic] Per-user Vertex SA lookup failed; falling back to global service key',
|
||||
err,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!userServiceKeyLoaded) {
|
||||
// Load credentials with optional YAML config overrides
|
||||
const credentialOptions = vertexConfig ? getVertexCredentialOptions(vertexConfig) : undefined;
|
||||
credentials = await loadAnthropicVertexCredentials(credentialOptions);
|
||||
}
|
||||
|
||||
// Store vertex options for client creation
|
||||
if (vertexConfig) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue