mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-23 02:36:12 +01:00
✨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:
parent
71c30a3640
commit
d99a9db3f6
58 changed files with 2146 additions and 1223 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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', () => {
|
||||
|
|
|
|||
|
|
@ -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.' };
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
|
|
|
|||
|
|
@ -932,6 +932,10 @@ export enum CacheKeys {
|
|||
* Key for in-progress messages.
|
||||
*/
|
||||
MESSAGES = 'messages',
|
||||
/**
|
||||
* Key for in-progress flow states.
|
||||
*/
|
||||
FLOWS = 'flows',
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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(() => ({}));
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
@ -73,5 +73,8 @@
|
|||
"diff": "^7.0.0",
|
||||
"eventsource": "^3.0.1",
|
||||
"express": "^4.21.2"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"keyv": "^4.5.4"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
152
packages/mcp/src/flow/manager.spec.ts
Normal file
152
packages/mcp/src/flow/manager.spec.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
241
packages/mcp/src/flow/manager.ts
Normal file
241
packages/mcp/src/flow/manager.ts
Normal 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
23
packages/mcp/src/flow/types.ts
Normal file
23
packages/mcp/src/flow/types.ts
Normal 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;
|
||||
}
|
||||
|
|
@ -1,4 +1,7 @@
|
|||
/* MCP */
|
||||
export * from './manager';
|
||||
/* Flow */
|
||||
export * from './flow/manager';
|
||||
/* types */
|
||||
export type * from './types/mcp';
|
||||
export type * from './flow/types';
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@
|
|||
"isolatedModules": true,
|
||||
"noEmit": true,
|
||||
"sourceMap": true,
|
||||
"baseUrl": "." // This should be the root of your package
|
||||
"baseUrl": "."
|
||||
},
|
||||
"ts-node": {
|
||||
"experimentalSpecifierResolution": "node",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue