Merge branch 'main' into feat/Multitenant-login-OIDC

packages/data-provider/package.json

@@ -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",

packages/data-provider/specs/actions.spec.ts

@@ -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', () => {

packages/data-provider/src/actions.ts

@@ -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';
 
@@ -11,12 +17,13 @@ export type ParametersSchema = {
   type: string;
   properties: Record<string, Reference | Schema>;
   required: string[];
+  additionalProperties?: boolean;
 };
 
 export type OpenAPISchema = OpenAPIV3.SchemaObject &
   ParametersSchema & {
-    items?: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject;
-  };
+  items?: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject;
+};
 
 export type ApiKeyCredentials = {
   api_key: string;
@@ -36,8 +43,8 @@ export type Credentials = ApiKeyCredentials | OAuthCredentials;
 type MediaTypeObject =
   | undefined
   | {
-      [media: string]: OpenAPIV3.MediaTypeObject | undefined;
-    };
+  [media: string]: OpenAPIV3.MediaTypeObject | undefined;
+};
 
 type RequestBodyObject = Omit<OpenAPIV3.RequestBodyObject, 'content'> & {
   content: MediaTypeObject;
@@ -118,28 +125,40 @@ function openAPISchemaToZod(schema: OpenAPISchema): z.ZodTypeAny | undefined {
   return handler(schema);
 }
 
+/**
+ * Class representing a function signature.
+ */
 export class FunctionSignature {
   name: string;
   description: string;
   parameters: ParametersSchema;
+  strict: boolean;
 
-  constructor(name: string, description: string, parameters: ParametersSchema) {
+  constructor(name: string, description: string, parameters: ParametersSchema, strict?: boolean) {
     this.name = name;
     this.description = description;
     this.parameters = parameters;
+    this.strict = strict ?? false;
   }
 
   toObjectTool(): FunctionTool {
+    const parameters = {
+      ...this.parameters,
+      additionalProperties: this.strict ? false : undefined,
+    };
+
     return {
       type: Tools.function,
       function: {
         name: this.name,
         description: this.description,
-        parameters: this.parameters,
+        parameters,
+        ...(this.strict ? { strict: this.strict } : {}),
       },
     };
   }
 }
+
 class RequestConfig {
   constructor(
     readonly domain: string,
@@ -176,7 +195,7 @@ class RequestExecutor {
     return this;
   }
 
-  async setAuth(metadata: ActionMetadata) {
+  async setAuth(metadata: ActionMetadataRuntime) {
     if (!metadata.auth) {
       return this;
     }
@@ -199,6 +218,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 +251,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 +281,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') {
@@ -355,7 +377,9 @@ function sanitizeOperationId(input: string) {
   return input.replace(/[^a-zA-Z0-9_-]/g, '');
 }
 
-/** Function to convert OpenAPI spec to function signatures and request builders */
+/**
+ * Converts an OpenAPI spec to function signatures and request builders.
+ */
 export function openapiToFunction(
   openapiSpec: OpenAPIV3.Document,
   generateZodSchemas = false,
@@ -374,12 +398,15 @@ export function openapiToFunction(
     for (const [method, operation] of Object.entries(methods as OpenAPIV3.PathsObject)) {
       const operationObj = operation as OpenAPIV3.OperationObject & {
         'x-openai-isConsequential'?: boolean;
+      } & {
+        'x-strict'?: boolean
       };
 
       // Operation ID is used as the function name
       const defaultOperationId = `${method}_${path}`;
       const operationId = operationObj.operationId || sanitizeOperationId(defaultOperationId);
       const description = operationObj.summary || operationObj.description || '';
+      const isStrict = operationObj['x-strict'] ?? false;
 
       const parametersSchema: OpenAPISchema = {
         type: 'object',
@@ -419,7 +446,7 @@ export function openapiToFunction(
         }
       }
 
-      const functionSignature = new FunctionSignature(operationId, description, parametersSchema);
+      const functionSignature = new FunctionSignature(operationId, description, parametersSchema, isStrict);
       functionSignatures.push(functionSignature);
 
       const actionRequest = new ActionRequest(
@@ -451,6 +478,9 @@ export type ValidationResult = {
   spec?: OpenAPIV3.Document;
 };
 
+/**
+ * Validates and parses an OpenAPI spec.
+ */
 export function validateAndParseOpenAPISpec(specString: string): ValidationResult {
   try {
     let parsedSpec;
@@ -511,6 +541,7 @@ export function validateAndParseOpenAPISpec(specString: string): ValidationResul
       spec: parsedSpec,
     };
   } catch (error) {
+    console.error(error);
     return { status: false, message: 'Error parsing OpenAPI spec.' };
   }
-}
+}

packages/data-provider/src/bedrock.ts

@@ -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>;

packages/data-provider/src/config.ts

@@ -1,4 +1,3 @@
-/* eslint-disable max-len */
 import { z } from 'zod';
 import type { ZodError } from 'zod';
 import type { TModelsConfig } from './types';
@@ -43,9 +42,8 @@ export const fileSourceSchema = z.nativeEnum(FileSources);
 type SchemaShape<T> = T extends z.ZodObject<infer U> ? U : never;
 
 // Helper type to determine the default value or undefined based on whether the field has a default
-type DefaultValue<T> = T extends z.ZodDefault<z.ZodTypeAny>
-  ? ReturnType<T['_def']['defaultValue']>
-  : undefined;
+type DefaultValue<T> =
+  T extends z.ZodDefault<z.ZodTypeAny> ? ReturnType<T['_def']['defaultValue']> : undefined;
 
 // Extract default values or undefined from the schema shape
 type ExtractDefaults<T> = {
@@ -145,6 +143,7 @@ export enum AgentCapabilities {
   end_after_tools = 'end_after_tools',
   execute_code = 'execute_code',
   file_search = 'file_search',
+  artifacts = 'artifacts',
   actions = 'actions',
   tools = 'tools',
 }
@@ -218,6 +217,7 @@ export const agentsEndpointSChema = baseEndpointSchema.merge(
       .default([
         AgentCapabilities.execute_code,
         AgentCapabilities.file_search,
+        AgentCapabilities.artifacts,
         AgentCapabilities.actions,
         AgentCapabilities.tools,
       ]),
@@ -951,6 +951,10 @@ export enum CacheKeys {
    * Key for in-progress messages.
    */
   MESSAGES = 'messages',
+  /**
+   * Key for in-progress flow states.
+   */
+  FLOWS = 'flows',
 }
 
 /**
@@ -1082,6 +1086,7 @@ export enum ImageDetailCost {
   /**
    * Additional Cost added to High Resolution Total Cost
    */
+  // eslint-disable-next-line @typescript-eslint/no-duplicate-enum-values
   ADDITIONAL = 85,
 }
 
@@ -1152,7 +1157,7 @@ export enum TTSProviders {
 /** Enum for app-wide constants */
 export enum Constants {
   /** Key for the app's version. */
-  VERSION = 'v0.7.6',
+  VERSION = 'v0.7.7-rc1',
   /** Key for the Custom Config's version (librechat.yaml). */
   CONFIG_VERSION = '1.2.1',
   /** Standard value for the first message's `parentMessageId` value, to indicate no parent exists. */

packages/data-provider/src/schemas.ts

@@ -155,6 +155,7 @@ export const defaultAgentFormValues = {
   tools: [],
   provider: {},
   projectIds: [],
+  artifacts: '',
   isCollaborative: false,
   [Tools.execute_code]: false,
   [Tools.file_search]: false,
@@ -877,7 +878,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 +889,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 +942,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 +1144,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 +1166,5 @@ export const compactAgentsSchema = tConversationSchema
     instructions: true,
     additional_instructions: true,
   })
-  .transform(removeNullishValues)
+  .transform((obj) => removeNullishValues(obj))
   .catch(() => ({}));

packages/data-provider/src/types/agents.ts

@@ -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 {

packages/data-provider/src/types/assistants.ts

@@ -3,6 +3,7 @@ import type { AssistantsEndpoint, AgentProvider } from 'src/schemas';
 import type { ContentTypes } from './runs';
 import type { Agents } from './agents';
 import type { TFile } from './files';
+import { ArtifactModes } from 'src/artifacts';
 
 export type Schema = OpenAPIV3.SchemaObject & { description?: string };
 export type Reference = OpenAPIV3.ReferenceObject & { description?: string };
@@ -38,6 +39,8 @@ export type FunctionTool = {
     description: string;
     name: string;
     parameters: Record<string, unknown>;
+    strict?: boolean;
+    additionalProperties?: boolean; // must be false if strict is true https://platform.openai.com/docs/guides/structured-outputs/some-type-specific-keywords-are-not-yet-supported
   };
 };
 
@@ -202,6 +205,7 @@ export type Agent = {
   created_at: number;
   avatar: AgentAvatar | null;
   instructions: string | null;
+  additional_instructions?: string | null;
   tools?: string[];
   projectIds?: string[];
   tool_kwargs?: Record<string, unknown>;
@@ -215,6 +219,7 @@ export type Agent = {
   agent_ids?: string[];
   end_after_tools?: boolean;
   hide_sequential_outputs?: boolean;
+  artifacts?: ArtifactModes;
 };
 
 export type TAgentsMap = Record<string, Agent | undefined>;
@@ -229,7 +234,7 @@ export type AgentCreateParams = {
   provider: AgentProvider;
   model: string | null;
   model_parameters: AgentModelParameters;
-} & Pick<Agent, 'agent_ids' | 'end_after_tools' | 'hide_sequential_outputs'>;
+} & Pick<Agent, 'agent_ids' | 'end_after_tools' | 'hide_sequential_outputs' | 'artifacts'>;
 
 export type AgentUpdateParams = {
   name?: string | null;
@@ -245,7 +250,7 @@ export type AgentUpdateParams = {
   projectIds?: string[];
   removeProjectIds?: string[];
   isCollaborative?: boolean;
-} & Pick<Agent, 'agent_ids' | 'end_after_tools' | 'hide_sequential_outputs'>;
+} & Pick<Agent, 'agent_ids' | 'end_after_tools' | 'hide_sequential_outputs' | 'artifacts'>;
 
 export type AgentListParams = {
   limit?: number;
@@ -417,6 +422,8 @@ export type PartMetadata = {
   asset_pointer?: string;
   status?: string;
   action?: boolean;
+  auth?: string;
+  expires_at?: number;
 };
 
 export type ContentPart = (
@@ -506,6 +513,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 = {

packages/data-provider/tsconfig.json

@@ -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"]
     }
   },

packages/mcp/jest.config.mjs

@@ -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,
+};

packages/mcp/package.json

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

packages/mcp/src/flow/manager.spec.ts

@@ -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);
+  });
+});

packages/mcp/src/flow/manager.ts

@@ -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;
+    }
+  }
+}

packages/mcp/src/flow/types.ts

@@ -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;
+}

packages/mcp/src/index.ts

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

packages/mcp/tsconfig.json

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