feat: OAuth for Actions (#5693)

* feat: OAuth for Actions

* WIP: PoC flow state manager

* refactor: Add identifier field to token model from action schema

* chore: fix potential file type issues

* ci: fix type issue with action metadata auth

* fix: ensure FlowManagerOptions has a default ttl value

* WIP: OAUTH actions

* WIP: first pass OAuth Action

* fix: standardize identifier usage in OAuth flow handling

* fix: update token retrieval to include userId in query and use correct identifier

* refacotr: update token retrieval to use userId for OAuth token query

* feat: Tool Call Auth styling

* fix: streamline token creation and add type field to token schema

* refactor: cleanup OAuth flow by encrypting client credentials and ensuring oauth operations only run under condition

* refactor: use encrypted credentials in OAuth callback

* fix: update Token collection indexes to use expiresAt TTL index and not createdAt legacy index

* refactor: enhance Token index cleanup by improving logging and removing redundant index creation logic

* refactor: remove unused OAuth login route and related logic for improved clarity

* refactor: replace fetch with axios for OAuth token exchange and improve error handling

* refactor: better UX after authentication before oauth tool execution

* refactor: implement cleanup handlers for FlowStateManager intervals to enhance resource management

* refactor: encrypt OAuth tokens before storing and decrypt upon retrieval for enhanced security

* refactor: enhance authentication success page with improved styling and countdown feature

* refactor: add response_type parameter to OAuth redirect URI for improved compatibility

* chore: update translation.json new localizations

* chore: remove unused OGDialog import from OGDialogTemplate component

* refactor: Actions Auth using new Dialog styling, use same component with Agents/Assistants

* refactor: update removeNullishValues function to support removal of empty strings and adjust transform usage in schemas

* chore: bump version of librechat-data-provider to 0.7.6991

* refactor: integrate removeNullishValues function to clean metadata before encryption in agent and assistant routes

* refactor: update OAuth input fields to use 'password' type for better security

* refactor: update localization placeholders for sign-in message to use double curly braces

* refactor: add access_type parameter for offline access in createActionTool function

* refactor: implement handleOAuthToken function for token management and encryption

* feat: refresh token support

* refactor: add default expiration for access token and error handling for missing token

* feat: localizations for ActionAuth

* refactor: set refresh token expiration to null to not expire if expiry never given

* fix: prevent crash fromerror within async handleAbortError in AskController, EditController, and AgentController

* feat: Action Callback URL

* 🌍 i18n: Update translation.json with latest translations

* refactor: handle errors in flow state checking to prevent unhandled promise rejections

* fix: improve flow state concurrency to prevent multiple token creation calls

* refactor: RequestExecutor to use separate axios instance

* refactor: improve concurrency flows by keeping completed state until TTL expiry

* refactor: increase TTL for flow state management and adjust monitoring interval

* ci: mock axios instance creation in actions spec

* feat: add Babel and Jest configuration files; implement FlowStateManager tests with concurrency handling

* chore: add disableOAuth prop to ActionsAuth (not implemented for Assistants yet)

---------

Co-authored-by: Danny Avila <danny@librechat.ai>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Ruben Talstra 2025-02-10 21:56:08 +01:00 committed by GitHub
parent 71c30a3640
commit d99a9db3f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
58 changed files with 2146 additions and 1223 deletions

View file

@ -1,6 +1,6 @@
{
"name": "librechat-data-provider",
"version": "0.7.699",
"version": "0.7.6991",
"description": "data services for librechat apps",
"main": "dist/index.js",
"module": "dist/index.es.js",

View file

@ -21,6 +21,7 @@ import type { ParametersSchema } from '../src/actions';
jest.mock('axios');
const mockedAxios = axios as jest.Mocked<typeof axios>;
mockedAxios.create.mockReturnValue(mockedAxios);
describe('FunctionSignature', () => {
it('creates a function signature and converts to JSON tool', () => {

View file

@ -1,9 +1,15 @@
import { z } from 'zod';
import axios from 'axios';
import _axios from 'axios';
import { URL } from 'url';
import crypto from 'crypto';
import { load } from 'js-yaml';
import type { FunctionTool, Schema, Reference, ActionMetadata } from './types/assistants';
import type {
FunctionTool,
Schema,
Reference,
ActionMetadata,
ActionMetadataRuntime,
} from './types/assistants';
import type { OpenAPIV3 } from 'openapi-types';
import { Tools, AuthTypeEnum, AuthorizationTypeEnum } from './types/assistants';
@ -176,7 +182,7 @@ class RequestExecutor {
return this;
}
async setAuth(metadata: ActionMetadata) {
async setAuth(metadata: ActionMetadataRuntime) {
if (!metadata.auth) {
return this;
}
@ -199,6 +205,8 @@ class RequestExecutor {
/* OAuth */
oauth_client_id,
oauth_client_secret,
oauth_token_expires_at,
oauth_access_token = '',
} = metadata;
const isApiKey = api_key != null && api_key.length > 0 && type === AuthTypeEnum.ServiceHttp;
@ -230,22 +238,23 @@ class RequestExecutor {
) {
this.authHeaders[custom_auth_header] = api_key;
} else if (isOAuth) {
const authToken = this.authToken ?? '';
if (!authToken) {
const tokenResponse = await axios.post(
client_url,
{
client_id: oauth_client_id,
client_secret: oauth_client_secret,
scope: scope,
grant_type: 'client_credentials',
},
{
headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
},
);
this.authToken = tokenResponse.data.access_token;
// TODO: maybe doing it in a different way later on. but we want that the user needs to folllow the oauth flow.
// If we do not have a valid token, bail or ask user to sign in
const now = new Date();
// 1. Check if token is set
if (!oauth_access_token) {
throw new Error('No access token found. Please log in first.');
}
// 2. Check if token is expired
if (oauth_token_expires_at && now >= new Date(oauth_token_expires_at)) {
// Optionally check refresh_token logic, or just prompt user to re-login
throw new Error('Access token is expired. Please re-login.');
}
// If valid, use it
this.authToken = oauth_access_token;
this.authHeaders['Authorization'] = `Bearer ${this.authToken}`;
}
return this;
@ -259,7 +268,7 @@ class RequestExecutor {
};
const method = this.config.method.toLowerCase();
const axios = _axios.create();
if (method === 'get') {
return axios.get(url, { headers, params: this.params });
} else if (method === 'post') {
@ -511,6 +520,7 @@ export function validateAndParseOpenAPISpec(specString: string): ValidationResul
spec: parsedSpec,
};
} catch (error) {
console.error(error);
return { status: false, message: 'Error parsing OpenAPI spec.' };
}
}

View file

@ -25,7 +25,7 @@ export const bedrockInputSchema = s.tConversationSchema
topK: true,
additionalModelRequestFields: true,
})
.transform(s.removeNullishValues)
.transform((obj) => s.removeNullishValues(obj))
.catch(() => ({}));
export type BedrockConverseInput = z.infer<typeof bedrockInputSchema>;

View file

@ -932,6 +932,10 @@ export enum CacheKeys {
* Key for in-progress messages.
*/
MESSAGES = 'messages',
/**
* Key for in-progress flow states.
*/
FLOWS = 'flows',
}
/**

View file

@ -877,7 +877,10 @@ export const gptPluginsSchema = tConversationSchema
maxContextTokens: undefined,
}));
export function removeNullishValues<T extends Record<string, unknown>>(obj: T): Partial<T> {
export function removeNullishValues<T extends Record<string, unknown>>(
obj: T,
removeEmptyStrings?: boolean,
): Partial<T> {
const newObj: Partial<T> = { ...obj };
(Object.keys(newObj) as Array<keyof T>).forEach((key) => {
@ -885,6 +888,9 @@ export function removeNullishValues<T extends Record<string, unknown>>(obj: T):
if (value === undefined || value === null) {
delete newObj[key];
}
if (removeEmptyStrings && typeof value === 'string' && value === '') {
delete newObj[key];
}
});
return newObj;
@ -935,8 +941,7 @@ export const compactAssistantSchema = tConversationSchema
greeting: true,
spec: true,
})
// will change after adding temperature
.transform(removeNullishValues)
.transform((obj) => removeNullishValues(obj))
.catch(() => ({}));
export const agentsSchema = tConversationSchema
@ -1138,7 +1143,7 @@ export const compactPluginsSchema = tConversationSchema
})
.catch(() => ({}));
const tBannerSchema = z.object({
export const tBannerSchema = z.object({
bannerId: z.string(),
message: z.string(),
displayFrom: z.string(),
@ -1160,5 +1165,5 @@ export const compactAgentsSchema = tConversationSchema
instructions: true,
additional_instructions: true,
})
.transform(removeNullishValues)
.transform((obj) => removeNullishValues(obj))
.catch(() => ({}));

View file

@ -52,6 +52,10 @@ export namespace Agents {
id?: string;
/** If provided, the output of the tool call */
output?: string;
/** Auth URL */
auth?: string;
/** Expiration time */
expires_at?: number;
};
export type ToolEndEvent = {
@ -190,6 +194,8 @@ export namespace Agents {
export type ToolCallDelta = {
type: StepTypes.TOOL_CALLS | string;
tool_calls?: ToolCallChunk[];
auth?: string;
expires_at?: number;
};
export type AgentToolCall = FunctionToolCall | ToolCall;
export interface ExtendedMessageContent {

View file

@ -417,6 +417,8 @@ export type PartMetadata = {
asset_pointer?: string;
status?: string;
action?: boolean;
auth?: string;
expires_at?: number;
};
export type ContentPart = (
@ -506,6 +508,12 @@ export type ActionMetadata = {
oauth_client_secret?: string;
};
export type ActionMetadataRuntime = ActionMetadata & {
oauth_access_token?: string;
oauth_refresh_token?: string;
oauth_token_expires_at?: Date;
};
/* Assistant types */
export type Action = {

View file

@ -18,9 +18,8 @@
"isolatedModules": true,
"noEmit": true,
"sourceMap": true,
"baseUrl": ".", // This should be the root of your package
"baseUrl": ".",
"paths": {
// Add path mappings
"librechat-data-provider/react-query": ["./src/react-query/index.ts"]
}
},

View file

@ -1,4 +1,4 @@
module.exports = {
export default {
collectCoverageFrom: ['src/**/*.{js,jsx,ts,tsx}', '!<rootDir>/node_modules/'],
coveragePathIgnorePatterns: ['/node_modules/', '/dist/'],
coverageReporters: ['text', 'cobertura'],
@ -15,4 +15,5 @@ module.exports = {
// },
// },
restoreMocks: true,
};
testTimeout: 15000,
};

View file

@ -73,5 +73,8 @@
"diff": "^7.0.0",
"eventsource": "^3.0.1",
"express": "^4.21.2"
},
"peerDependencies": {
"keyv": "^4.5.4"
}
}

View file

@ -0,0 +1,152 @@
import { FlowStateManager } from './manager';
import Keyv from 'keyv';
import type { FlowState } from './types';
// Create a mock class without extending Keyv
class MockKeyv {
private store: Map<string, FlowState<string>>;
constructor() {
this.store = new Map();
}
async get(key: string): Promise<FlowState<string> | undefined> {
return this.store.get(key);
}
// eslint-disable-next-line @typescript-eslint/no-unused-vars
async set(key: string, value: FlowState<string>, _ttl?: number): Promise<true> {
this.store.set(key, value);
return true;
}
async delete(key: string): Promise<boolean> {
return this.store.delete(key);
}
}
describe('FlowStateManager', () => {
let flowManager: FlowStateManager<string>;
let store: MockKeyv;
beforeEach(() => {
store = new MockKeyv();
// Type assertion here since we know our mock implements the necessary methods
flowManager = new FlowStateManager(store as unknown as Keyv, { ttl: 30000, ci: true });
});
afterEach(() => {
jest.clearAllMocks();
});
describe('Concurrency Tests', () => {
it('should handle concurrent flow creation and return same result', async () => {
const flowId = 'test-flow';
const type = 'test-type';
// Start two concurrent flow creations
const flow1Promise = flowManager.createFlowWithHandler(flowId, type, async () => {
await new Promise((resolve) => setTimeout(resolve, 100));
return 'result';
});
const flow2Promise = flowManager.createFlowWithHandler(flowId, type, async () => {
await new Promise((resolve) => setTimeout(resolve, 50));
return 'different-result';
});
// Both should resolve to the same result from the first handler
const [result1, result2] = await Promise.all([flow1Promise, flow2Promise]);
expect(result1).toBe('result');
expect(result2).toBe('result');
});
it('should handle flow timeout correctly', async () => {
const flowId = 'timeout-flow';
const type = 'test-type';
// Create flow with very short TTL
const shortTtlManager = new FlowStateManager(store as unknown as Keyv, {
ttl: 100,
ci: true,
});
const flowPromise = shortTtlManager.createFlow(flowId, type);
await expect(flowPromise).rejects.toThrow('test-type flow timed out');
});
it('should maintain flow state consistency under high concurrency', async () => {
const flowId = 'concurrent-flow';
const type = 'test-type';
// Create multiple concurrent operations
const operations = [];
for (let i = 0; i < 10; i++) {
operations.push(
flowManager.createFlowWithHandler(flowId, type, async () => {
await new Promise((resolve) => setTimeout(resolve, Math.random() * 50));
return `result-${i}`;
}),
);
}
// All operations should resolve to the same result
const results = await Promise.all(operations);
const firstResult = results[0];
results.forEach((result: string) => {
expect(result).toBe(firstResult);
});
});
it('should handle race conditions in flow completion', async () => {
const flowId = 'test-flow';
const type = 'test-type';
// Create initial flow
const flowPromise = flowManager.createFlow(flowId, type);
// Increase delay to ensure flow is properly created
await new Promise((resolve) => setTimeout(resolve, 500));
// Complete the flow
await flowManager.completeFlow(flowId, type, 'result1');
const result = await flowPromise;
expect(result).toBe('result1');
}, 15000);
it('should handle concurrent flow monitoring', async () => {
const flowId = 'test-flow';
const type = 'test-type';
// Create initial flow
const flowPromise = flowManager.createFlow(flowId, type);
// Increase delay
await new Promise((resolve) => setTimeout(resolve, 500));
// Complete the flow
await flowManager.completeFlow(flowId, type, 'success');
const result = await flowPromise;
expect(result).toBe('success');
}, 15000);
it('should handle concurrent success and failure attempts', async () => {
const flowId = 'race-flow';
const type = 'test-type';
const flowPromise = flowManager.createFlow(flowId, type);
// Increase delay
await new Promise((resolve) => setTimeout(resolve, 500));
// Fail the flow
await flowManager.failFlow(flowId, type, new Error('failure'));
await expect(flowPromise).rejects.toThrow('failure');
}, 15000);
});
});

View file

@ -0,0 +1,241 @@
import Keyv from 'keyv';
import type { Logger } from 'winston';
import type { FlowState, FlowMetadata, FlowManagerOptions } from './types';
export class FlowStateManager<T = unknown> {
private keyv: Keyv;
private ttl: number;
private logger: Logger;
private intervals: Set<NodeJS.Timeout>;
private static getDefaultLogger(): Logger {
return {
error: console.error,
warn: console.warn,
info: console.info,
debug: console.debug,
} as Logger;
}
constructor(store: Keyv, options?: FlowManagerOptions) {
if (!options) {
options = { ttl: 60000 * 3 };
}
const { ci = false, ttl, logger } = options;
if (!ci && !(store instanceof Keyv)) {
throw new Error('Invalid store provided to FlowStateManager');
}
this.ttl = ttl;
this.keyv = store;
this.logger = logger || FlowStateManager.getDefaultLogger();
this.intervals = new Set();
this.setupCleanupHandlers();
}
private setupCleanupHandlers() {
const cleanup = () => {
this.logger.info('Cleaning up FlowStateManager intervals...');
this.intervals.forEach((interval) => clearInterval(interval));
this.intervals.clear();
process.exit(0);
};
process.on('SIGTERM', cleanup);
process.on('SIGINT', cleanup);
process.on('SIGQUIT', cleanup);
process.on('SIGHUP', cleanup);
}
private getFlowKey(flowId: string, type: string): string {
return `${type}:${flowId}`;
}
/**
* Creates a new flow and waits for its completion
*/
async createFlow(flowId: string, type: string, metadata: FlowMetadata = {}): Promise<T> {
const flowKey = this.getFlowKey(flowId, type);
let existingState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (existingState) {
this.logger.debug(`[${flowKey}] Flow already exists`);
return this.monitorFlow(flowKey, type);
}
await new Promise((resolve) => setTimeout(resolve, 250));
existingState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (existingState) {
this.logger.debug(`[${flowKey}] Flow exists on 2nd check`);
return this.monitorFlow(flowKey, type);
}
const initialState: FlowState = {
type,
status: 'PENDING',
metadata,
createdAt: Date.now(),
};
this.logger.debug('Creating initial flow state:', flowKey);
await this.keyv.set(flowKey, initialState, this.ttl);
return this.monitorFlow(flowKey, type);
}
private monitorFlow(flowKey: string, type: string): Promise<T> {
return new Promise<T>((resolve, reject) => {
const checkInterval = 2000;
let elapsedTime = 0;
const intervalId = setInterval(async () => {
try {
const flowState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (!flowState) {
clearInterval(intervalId);
this.intervals.delete(intervalId);
this.logger.error(`[${flowKey}] Flow state not found`);
reject(new Error(`${type} Flow state not found`));
return;
}
if (flowState.status !== 'PENDING') {
clearInterval(intervalId);
this.intervals.delete(intervalId);
this.logger.debug(`[${flowKey}] Flow completed`);
if (flowState.status === 'COMPLETED' && flowState.result !== undefined) {
resolve(flowState.result);
} else if (flowState.status === 'FAILED') {
await this.keyv.delete(flowKey);
reject(new Error(flowState.error ?? `${type} flow failed`));
}
return;
}
elapsedTime += checkInterval;
if (elapsedTime >= this.ttl) {
clearInterval(intervalId);
this.intervals.delete(intervalId);
this.logger.error(
`[${flowKey}] Flow timed out | Elapsed time: ${elapsedTime} | TTL: ${this.ttl}`,
);
await this.keyv.delete(flowKey);
reject(new Error(`${type} flow timed out`));
}
this.logger.debug(
`[${flowKey}] Flow state elapsed time: ${elapsedTime}, checking again...`,
);
} catch (error) {
this.logger.error(`[${flowKey}] Error checking flow state:`, error);
clearInterval(intervalId);
this.intervals.delete(intervalId);
reject(error);
}
}, checkInterval);
this.intervals.add(intervalId);
});
}
/**
* Completes a flow successfully
*/
async completeFlow(flowId: string, type: string, result: T): Promise<boolean> {
const flowKey = this.getFlowKey(flowId, type);
const flowState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (!flowState) {
return false;
}
const updatedState: FlowState<T> = {
...flowState,
status: 'COMPLETED',
result,
completedAt: Date.now(),
};
await this.keyv.set(flowKey, updatedState, this.ttl);
return true;
}
/**
* Marks a flow as failed
*/
async failFlow(flowId: string, type: string, error: Error | string): Promise<boolean> {
const flowKey = this.getFlowKey(flowId, type);
const flowState = (await this.keyv.get(flowKey)) as FlowState | undefined;
if (!flowState) {
return false;
}
const updatedState: FlowState = {
...flowState,
status: 'FAILED',
error: error instanceof Error ? error.message : error,
failedAt: Date.now(),
};
await this.keyv.set(flowKey, updatedState, this.ttl);
return true;
}
/**
* Gets current flow state
*/
async getFlowState(flowId: string, type: string): Promise<FlowState<T> | null> {
const flowKey = this.getFlowKey(flowId, type);
return this.keyv.get(flowKey);
}
/**
* Creates a new flow and waits for its completion, only executing the handler if no existing flow is found
* @param flowId - The ID of the flow
* @param type - The type of flow
* @param handler - Async function to execute if no existing flow is found
* @param metadata - Optional metadata for the flow
*/
async createFlowWithHandler(
flowId: string,
type: string,
handler: () => Promise<T>,
metadata: FlowMetadata = {},
): Promise<T> {
const flowKey = this.getFlowKey(flowId, type);
let existingState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (existingState) {
this.logger.debug(`[${flowKey}] Flow already exists`);
return this.monitorFlow(flowKey, type);
}
await new Promise((resolve) => setTimeout(resolve, 250));
existingState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (existingState) {
this.logger.debug(`[${flowKey}] Flow exists on 2nd check`);
return this.monitorFlow(flowKey, type);
}
const initialState: FlowState = {
type,
status: 'PENDING',
metadata,
createdAt: Date.now(),
};
this.logger.debug(`[${flowKey}] Creating initial flow state`);
await this.keyv.set(flowKey, initialState, this.ttl);
try {
const result = await handler();
await this.completeFlow(flowId, type, result);
return result;
} catch (error) {
await this.failFlow(flowId, type, error instanceof Error ? error : new Error(String(error)));
throw error;
}
}
}

View file

@ -0,0 +1,23 @@
import type { Logger } from 'winston';
export type FlowStatus = 'PENDING' | 'COMPLETED' | 'FAILED';
export interface FlowMetadata {
[key: string]: unknown;
}
export interface FlowState<T = unknown> {
type: string;
status: FlowStatus;
metadata: FlowMetadata;
createdAt: number;
result?: T;
error?: string;
completedAt?: number;
failedAt?: number;
}
export interface FlowManagerOptions {
ttl: number;
ci?: boolean;
logger?: Logger;
}

View file

@ -1,4 +1,7 @@
/* MCP */
export * from './manager';
/* Flow */
export * from './flow/manager';
/* types */
export type * from './types/mcp';
export type * from './flow/types';

View file

@ -18,7 +18,7 @@
"isolatedModules": true,
"noEmit": true,
"sourceMap": true,
"baseUrl": "." // This should be the root of your package
"baseUrl": "."
},
"ts-node": {
"experimentalSpecifierResolution": "node",