mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-04-07 00:15:23 +02:00
Merge branch 'main' into dt-conf-logo
This commit is contained in:
commit
74c8f518c2
455 changed files with 46339 additions and 7322 deletions
25
.env.example
25
.env.example
|
|
@ -47,6 +47,10 @@ TRUST_PROXY=1
|
|||
# password policies.
|
||||
# MIN_PASSWORD_LENGTH=8
|
||||
|
||||
# When enabled, the app will continue running after encountering uncaught exceptions
|
||||
# instead of exiting the process. Not recommended for production unless necessary.
|
||||
# CONTINUE_ON_UNCAUGHT_EXCEPTION=false
|
||||
|
||||
#===============#
|
||||
# JSON Logging #
|
||||
#===============#
|
||||
|
|
@ -87,6 +91,16 @@ NODE_MAX_OLD_SPACE_SIZE=6144
|
|||
|
||||
# CONFIG_PATH="/alternative/path/to/librechat.yaml"
|
||||
|
||||
#==================#
|
||||
# Langfuse Tracing #
|
||||
#==================#
|
||||
|
||||
# Get Langfuse API keys for your project from the project settings page: https://cloud.langfuse.com
|
||||
|
||||
# LANGFUSE_PUBLIC_KEY=
|
||||
# LANGFUSE_SECRET_KEY=
|
||||
# LANGFUSE_BASE_URL=
|
||||
|
||||
#===================================================#
|
||||
# Endpoints #
|
||||
#===================================================#
|
||||
|
|
@ -121,7 +135,7 @@ PROXY=
|
|||
#============#
|
||||
|
||||
ANTHROPIC_API_KEY=user_provided
|
||||
# ANTHROPIC_MODELS=claude-opus-4-20250514,claude-sonnet-4-20250514,claude-3-7-sonnet-20250219,claude-3-5-sonnet-20241022,claude-3-5-haiku-20241022,claude-3-opus-20240229,claude-3-sonnet-20240229,claude-3-haiku-20240307
|
||||
# ANTHROPIC_MODELS=claude-sonnet-4-6,claude-opus-4-6,claude-opus-4-20250514,claude-sonnet-4-20250514,claude-3-7-sonnet-20250219,claude-3-5-sonnet-20241022,claude-3-5-haiku-20241022,claude-3-opus-20240229,claude-3-sonnet-20240229,claude-3-haiku-20240307
|
||||
# ANTHROPIC_REVERSE_PROXY=
|
||||
|
||||
# Set to true to use Anthropic models through Google Vertex AI instead of direct API
|
||||
|
|
@ -156,7 +170,8 @@ ANTHROPIC_API_KEY=user_provided
|
|||
# BEDROCK_AWS_SESSION_TOKEN=someSessionToken
|
||||
|
||||
# Note: This example list is not meant to be exhaustive. If omitted, all known, supported model IDs will be included for you.
|
||||
# BEDROCK_AWS_MODELS=anthropic.claude-3-5-sonnet-20240620-v1:0,meta.llama3-1-8b-instruct-v1:0
|
||||
# BEDROCK_AWS_MODELS=anthropic.claude-sonnet-4-6,anthropic.claude-opus-4-6-v1,anthropic.claude-3-5-sonnet-20240620-v1:0,meta.llama3-1-8b-instruct-v1:0
|
||||
# Cross-region inference model IDs: us.anthropic.claude-sonnet-4-6,us.anthropic.claude-opus-4-6-v1,global.anthropic.claude-opus-4-6-v1
|
||||
|
||||
# See all Bedrock model IDs here: https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html#model-ids-arns
|
||||
|
||||
|
|
@ -737,8 +752,10 @@ HELP_AND_FAQ_URL=https://librechat.ai
|
|||
# REDIS_PING_INTERVAL=300
|
||||
|
||||
# Force specific cache namespaces to use in-memory storage even when Redis is enabled
|
||||
# Comma-separated list of CacheKeys (e.g., ROLES,MESSAGES)
|
||||
# FORCED_IN_MEMORY_CACHE_NAMESPACES=ROLES,MESSAGES
|
||||
# Comma-separated list of CacheKeys
|
||||
# Defaults to CONFIG_STORE,APP_CONFIG so YAML-derived config stays per-container (safe for blue/green deployments)
|
||||
# Set to empty string to force all namespaces through Redis: FORCED_IN_MEMORY_CACHE_NAMESPACES=
|
||||
# FORCED_IN_MEMORY_CACHE_NAMESPACES=CONFIG_STORE,APP_CONFIG
|
||||
|
||||
# Leader Election Configuration (for multi-instance deployments with Redis)
|
||||
# Duration in seconds that the leader lease is valid before it expires (default: 25)
|
||||
|
|
|
|||
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -15,6 +15,7 @@ pids
|
|||
|
||||
# CI/CD data
|
||||
test-image*
|
||||
dump.rdb
|
||||
|
||||
# Directory for instrumented libs generated by jscoverage/JSCover
|
||||
lib-cov
|
||||
|
|
@ -29,6 +30,9 @@ coverage
|
|||
config/translations/stores/*
|
||||
client/src/localization/languages/*_missing_keys.json
|
||||
|
||||
# Turborepo
|
||||
.turbo
|
||||
|
||||
# Compiled Dirs (http://nodejs.org/api/addons.html)
|
||||
build/
|
||||
dist/
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
# v0.8.2-rc2
|
||||
# v0.8.2
|
||||
|
||||
# Base node image
|
||||
FROM node:20-alpine AS node
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
# Dockerfile.multi
|
||||
# v0.8.2-rc2
|
||||
# v0.8.2
|
||||
|
||||
# Set configurable max-old-space-size with default
|
||||
ARG NODE_MAX_OLD_SPACE_SIZE=6144
|
||||
|
|
|
|||
13
README.md
13
README.md
|
|
@ -109,6 +109,11 @@
|
|||
- 🎨 **Customizable Interface**:
|
||||
- Customizable Dropdown & Interface that adapts to both power users and newcomers
|
||||
|
||||
- 🌊 **[Resumable Streams](https://www.librechat.ai/docs/features/resumable_streams)**:
|
||||
- Never lose a response: AI responses automatically reconnect and resume if your connection drops
|
||||
- Multi-Tab & Multi-Device Sync: Open the same chat in multiple tabs or pick up on another device
|
||||
- Production-Ready: Works from single-server setups to horizontally scaled deployments with Redis
|
||||
|
||||
- 🗣️ **Speech & Audio**:
|
||||
- Chat hands-free with Speech-to-Text and Text-to-Speech
|
||||
- Automatically send and play Audio
|
||||
|
|
@ -137,13 +142,11 @@
|
|||
|
||||
## 🪶 All-In-One AI Conversations with LibreChat
|
||||
|
||||
LibreChat brings together the future of assistant AIs with the revolutionary technology of OpenAI's ChatGPT. Celebrating the original styling, LibreChat gives you the ability to integrate multiple AI models. It also integrates and enhances original client features such as conversation and message search, prompt templates and plugins.
|
||||
LibreChat is a self-hosted AI chat platform that unifies all major AI providers in a single, privacy-focused interface.
|
||||
|
||||
With LibreChat, you no longer need to opt for ChatGPT Plus and can instead use free or pay-per-call APIs. We welcome contributions, cloning, and forking to enhance the capabilities of this advanced chatbot platform.
|
||||
Beyond chat, LibreChat provides AI Agents, Model Context Protocol (MCP) support, Artifacts, Code Interpreter, custom actions, conversation search, and enterprise-ready multi-user authentication.
|
||||
|
||||
[](https://www.youtube.com/watch?v=ilfwGQtJNlI)
|
||||
|
||||
Click on the thumbnail to open the video☝️
|
||||
Open source, actively developed, and built for anyone who values control over their AI infrastructure.
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -41,9 +41,9 @@ jest.mock('~/models', () => ({
|
|||
const { getConvo, saveConvo } = require('~/models');
|
||||
|
||||
jest.mock('@librechat/agents', () => {
|
||||
const { Providers } = jest.requireActual('@librechat/agents');
|
||||
const actual = jest.requireActual('@librechat/agents');
|
||||
return {
|
||||
Providers,
|
||||
...actual,
|
||||
ChatOpenAI: jest.fn().mockImplementation(() => {
|
||||
return {};
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -57,19 +57,6 @@
|
|||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Browser",
|
||||
"pluginKey": "web-browser",
|
||||
"description": "Scrape and summarize webpage data",
|
||||
"icon": "assets/web-browser.svg",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "OPENAI_API_KEY",
|
||||
"label": "OpenAI API Key",
|
||||
"description": "Browser makes use of OpenAI embeddings"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "DALL-E-3",
|
||||
"pluginKey": "dalle",
|
||||
|
|
|
|||
|
|
@ -1,14 +1,28 @@
|
|||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { SearchClient, AzureKeyCredential } = require('@azure/search-documents');
|
||||
|
||||
const azureAISearchJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
query: {
|
||||
type: 'string',
|
||||
description: 'Search word or phrase to Azure AI Search',
|
||||
},
|
||||
},
|
||||
required: ['query'],
|
||||
};
|
||||
|
||||
class AzureAISearch extends Tool {
|
||||
// Constants for default values
|
||||
static DEFAULT_API_VERSION = '2023-11-01';
|
||||
static DEFAULT_QUERY_TYPE = 'simple';
|
||||
static DEFAULT_TOP = 5;
|
||||
|
||||
static get jsonSchema() {
|
||||
return azureAISearchJsonSchema;
|
||||
}
|
||||
|
||||
// Helper function for initializing properties
|
||||
_initializeField(field, envVar, defaultValue) {
|
||||
return field || process.env[envVar] || defaultValue;
|
||||
|
|
@ -22,10 +36,7 @@ class AzureAISearch extends Tool {
|
|||
/* Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
|
||||
// Define schema
|
||||
this.schema = z.object({
|
||||
query: z.string().describe('Search word or phrase to Azure AI Search'),
|
||||
});
|
||||
this.schema = azureAISearchJsonSchema;
|
||||
|
||||
// Initialize properties using helper function
|
||||
this.serviceEndpoint = this._initializeField(
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
const { z } = require('zod');
|
||||
const path = require('path');
|
||||
const OpenAI = require('openai');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
|
|
@ -8,6 +7,36 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { getImageBasename, extractBaseURL } = require('@librechat/api');
|
||||
const { FileContext, ContentTypes } = require('librechat-data-provider');
|
||||
|
||||
const dalle3JsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
prompt: {
|
||||
type: 'string',
|
||||
maxLength: 4000,
|
||||
description:
|
||||
'A text description of the desired image, following the rules, up to 4000 characters.',
|
||||
},
|
||||
style: {
|
||||
type: 'string',
|
||||
enum: ['vivid', 'natural'],
|
||||
description:
|
||||
'Must be one of `vivid` or `natural`. `vivid` generates hyper-real and dramatic images, `natural` produces more natural, less hyper-real looking images',
|
||||
},
|
||||
quality: {
|
||||
type: 'string',
|
||||
enum: ['hd', 'standard'],
|
||||
description: 'The quality of the generated image. Only `hd` and `standard` are supported.',
|
||||
},
|
||||
size: {
|
||||
type: 'string',
|
||||
enum: ['1024x1024', '1792x1024', '1024x1792'],
|
||||
description:
|
||||
'The size of the requested image. Use 1024x1024 (square) as the default, 1792x1024 if the user requests a wide image, and 1024x1792 for full-body portraits. Always include this parameter in the request.',
|
||||
},
|
||||
},
|
||||
required: ['prompt', 'style', 'quality', 'size'],
|
||||
};
|
||||
|
||||
const displayMessage =
|
||||
"DALL-E displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
class DALLE3 extends Tool {
|
||||
|
|
@ -72,27 +101,11 @@ class DALLE3 extends Tool {
|
|||
// The prompt must intricately describe every part of the image in concrete, objective detail. THINK about what the end goal of the description is, and extrapolate that to what would make satisfying images.
|
||||
// All descriptions sent to dalle should be a paragraph of text that is extremely descriptive and detailed. Each should be more than 3 sentences long.
|
||||
// - The "vivid" style is HIGHLY preferred, but "natural" is also supported.`;
|
||||
this.schema = z.object({
|
||||
prompt: z
|
||||
.string()
|
||||
.max(4000)
|
||||
.describe(
|
||||
'A text description of the desired image, following the rules, up to 4000 characters.',
|
||||
),
|
||||
style: z
|
||||
.enum(['vivid', 'natural'])
|
||||
.describe(
|
||||
'Must be one of `vivid` or `natural`. `vivid` generates hyper-real and dramatic images, `natural` produces more natural, less hyper-real looking images',
|
||||
),
|
||||
quality: z
|
||||
.enum(['hd', 'standard'])
|
||||
.describe('The quality of the generated image. Only `hd` and `standard` are supported.'),
|
||||
size: z
|
||||
.enum(['1024x1024', '1792x1024', '1024x1792'])
|
||||
.describe(
|
||||
'The size of the requested image. Use 1024x1024 (square) as the default, 1792x1024 if the user requests a wide image, and 1024x1792 for full-body portraits. Always include this parameter in the request.',
|
||||
),
|
||||
});
|
||||
this.schema = dalle3JsonSchema;
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return dalle3JsonSchema;
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const fetch = require('node-fetch');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
|
|
@ -7,6 +6,84 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { FileContext, ContentTypes } = require('librechat-data-provider');
|
||||
|
||||
const fluxApiJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
action: {
|
||||
type: 'string',
|
||||
enum: ['generate', 'list_finetunes', 'generate_finetuned'],
|
||||
description:
|
||||
'Action to perform: "generate" for image generation, "generate_finetuned" for finetuned model generation, "list_finetunes" to get available custom models',
|
||||
},
|
||||
prompt: {
|
||||
type: 'string',
|
||||
description:
|
||||
'Text prompt for image generation. Required when action is "generate". Not used for list_finetunes.',
|
||||
},
|
||||
width: {
|
||||
type: 'number',
|
||||
description:
|
||||
'Width of the generated image in pixels. Must be a multiple of 32. Default is 1024.',
|
||||
},
|
||||
height: {
|
||||
type: 'number',
|
||||
description:
|
||||
'Height of the generated image in pixels. Must be a multiple of 32. Default is 768.',
|
||||
},
|
||||
prompt_upsampling: {
|
||||
type: 'boolean',
|
||||
description: 'Whether to perform upsampling on the prompt.',
|
||||
},
|
||||
steps: {
|
||||
type: 'integer',
|
||||
description: 'Number of steps to run the model for, a number from 1 to 50. Default is 40.',
|
||||
},
|
||||
seed: {
|
||||
type: 'number',
|
||||
description: 'Optional seed for reproducibility.',
|
||||
},
|
||||
safety_tolerance: {
|
||||
type: 'number',
|
||||
description:
|
||||
'Tolerance level for input and output moderation. Between 0 and 6, 0 being most strict, 6 being least strict.',
|
||||
},
|
||||
endpoint: {
|
||||
type: 'string',
|
||||
enum: [
|
||||
'/v1/flux-pro-1.1',
|
||||
'/v1/flux-pro',
|
||||
'/v1/flux-dev',
|
||||
'/v1/flux-pro-1.1-ultra',
|
||||
'/v1/flux-pro-finetuned',
|
||||
'/v1/flux-pro-1.1-ultra-finetuned',
|
||||
],
|
||||
description: 'Endpoint to use for image generation.',
|
||||
},
|
||||
raw: {
|
||||
type: 'boolean',
|
||||
description:
|
||||
'Generate less processed, more natural-looking images. Only works for /v1/flux-pro-1.1-ultra.',
|
||||
},
|
||||
finetune_id: {
|
||||
type: 'string',
|
||||
description: 'ID of the finetuned model to use',
|
||||
},
|
||||
finetune_strength: {
|
||||
type: 'number',
|
||||
description: 'Strength of the finetuning effect (typically between 0.1 and 1.2)',
|
||||
},
|
||||
guidance: {
|
||||
type: 'number',
|
||||
description: 'Guidance scale for finetuned models',
|
||||
},
|
||||
aspect_ratio: {
|
||||
type: 'string',
|
||||
description: 'Aspect ratio for ultra models (e.g., "16:9")',
|
||||
},
|
||||
},
|
||||
required: [],
|
||||
};
|
||||
|
||||
const displayMessage =
|
||||
"Flux displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
|
|
@ -57,82 +134,11 @@ class FluxAPI extends Tool {
|
|||
// Add base URL from environment variable with fallback
|
||||
this.baseUrl = process.env.FLUX_API_BASE_URL || 'https://api.us1.bfl.ai';
|
||||
|
||||
// Define the schema for structured input
|
||||
this.schema = z.object({
|
||||
action: z
|
||||
.enum(['generate', 'list_finetunes', 'generate_finetuned'])
|
||||
.default('generate')
|
||||
.describe(
|
||||
'Action to perform: "generate" for image generation, "generate_finetuned" for finetuned model generation, "list_finetunes" to get available custom models',
|
||||
),
|
||||
prompt: z
|
||||
.string()
|
||||
.optional()
|
||||
.describe(
|
||||
'Text prompt for image generation. Required when action is "generate". Not used for list_finetunes.',
|
||||
),
|
||||
width: z
|
||||
.number()
|
||||
.optional()
|
||||
.describe(
|
||||
'Width of the generated image in pixels. Must be a multiple of 32. Default is 1024.',
|
||||
),
|
||||
height: z
|
||||
.number()
|
||||
.optional()
|
||||
.describe(
|
||||
'Height of the generated image in pixels. Must be a multiple of 32. Default is 768.',
|
||||
),
|
||||
prompt_upsampling: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe('Whether to perform upsampling on the prompt.'),
|
||||
steps: z
|
||||
.number()
|
||||
.int()
|
||||
.optional()
|
||||
.describe('Number of steps to run the model for, a number from 1 to 50. Default is 40.'),
|
||||
seed: z.number().optional().describe('Optional seed for reproducibility.'),
|
||||
safety_tolerance: z
|
||||
.number()
|
||||
.optional()
|
||||
.default(6)
|
||||
.describe(
|
||||
'Tolerance level for input and output moderation. Between 0 and 6, 0 being most strict, 6 being least strict.',
|
||||
),
|
||||
endpoint: z
|
||||
.enum([
|
||||
'/v1/flux-pro-1.1',
|
||||
'/v1/flux-pro',
|
||||
'/v1/flux-dev',
|
||||
'/v1/flux-pro-1.1-ultra',
|
||||
'/v1/flux-pro-finetuned',
|
||||
'/v1/flux-pro-1.1-ultra-finetuned',
|
||||
])
|
||||
.optional()
|
||||
.default('/v1/flux-pro-1.1')
|
||||
.describe('Endpoint to use for image generation.'),
|
||||
raw: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe(
|
||||
'Generate less processed, more natural-looking images. Only works for /v1/flux-pro-1.1-ultra.',
|
||||
),
|
||||
finetune_id: z.string().optional().describe('ID of the finetuned model to use'),
|
||||
finetune_strength: z
|
||||
.number()
|
||||
.optional()
|
||||
.default(1.1)
|
||||
.describe('Strength of the finetuning effect (typically between 0.1 and 1.2)'),
|
||||
guidance: z.number().optional().default(2.5).describe('Guidance scale for finetuned models'),
|
||||
aspect_ratio: z
|
||||
.string()
|
||||
.optional()
|
||||
.default('16:9')
|
||||
.describe('Aspect ratio for ultra models (e.g., "16:9")'),
|
||||
});
|
||||
this.schema = fluxApiJsonSchema;
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return fluxApiJsonSchema;
|
||||
}
|
||||
|
||||
getAxiosConfig() {
|
||||
|
|
|
|||
|
|
@ -1,12 +1,33 @@
|
|||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
const googleSearchJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
query: {
|
||||
type: 'string',
|
||||
minLength: 1,
|
||||
description: 'The search query string.',
|
||||
},
|
||||
max_results: {
|
||||
type: 'integer',
|
||||
minimum: 1,
|
||||
maximum: 10,
|
||||
description: 'The maximum number of search results to return. Defaults to 5.',
|
||||
},
|
||||
},
|
||||
required: ['query'],
|
||||
};
|
||||
|
||||
class GoogleSearchResults extends Tool {
|
||||
static lc_name() {
|
||||
return 'google';
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return googleSearchJsonSchema;
|
||||
}
|
||||
|
||||
constructor(fields = {}) {
|
||||
super(fields);
|
||||
this.name = 'google';
|
||||
|
|
@ -28,25 +49,11 @@ class GoogleSearchResults extends Tool {
|
|||
this.description =
|
||||
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.';
|
||||
|
||||
this.schema = z.object({
|
||||
query: z.string().min(1).describe('The search query string.'),
|
||||
max_results: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(10)
|
||||
.optional()
|
||||
.describe('The maximum number of search results to return. Defaults to 10.'),
|
||||
// Note: Google API has its own parameters for search customization, adjust as needed.
|
||||
});
|
||||
this.schema = googleSearchJsonSchema;
|
||||
}
|
||||
|
||||
async _call(input) {
|
||||
const validationResult = this.schema.safeParse(input);
|
||||
if (!validationResult.success) {
|
||||
throw new Error(`Validation failed: ${JSON.stringify(validationResult.error.issues)}`);
|
||||
}
|
||||
|
||||
const { query, max_results = 5 } = validationResult.data;
|
||||
const { query, max_results = 5 } = input;
|
||||
|
||||
const response = await fetch(
|
||||
`https://www.googleapis.com/customsearch/v1?key=${this.apiKey}&cx=${
|
||||
|
|
|
|||
|
|
@ -1,8 +1,52 @@
|
|||
const { Tool } = require('@langchain/core/tools');
|
||||
const { z } = require('zod');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
const fetch = require('node-fetch');
|
||||
|
||||
const openWeatherJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
action: {
|
||||
type: 'string',
|
||||
enum: ['help', 'current_forecast', 'timestamp', 'daily_aggregation', 'overview'],
|
||||
description: 'The action to perform',
|
||||
},
|
||||
city: {
|
||||
type: 'string',
|
||||
description: 'City name for geocoding if lat/lon not provided',
|
||||
},
|
||||
lat: {
|
||||
type: 'number',
|
||||
description: 'Latitude coordinate',
|
||||
},
|
||||
lon: {
|
||||
type: 'number',
|
||||
description: 'Longitude coordinate',
|
||||
},
|
||||
exclude: {
|
||||
type: 'string',
|
||||
description: 'Parts to exclude from the response',
|
||||
},
|
||||
units: {
|
||||
type: 'string',
|
||||
enum: ['Celsius', 'Kelvin', 'Fahrenheit'],
|
||||
description: 'Temperature units',
|
||||
},
|
||||
lang: {
|
||||
type: 'string',
|
||||
description: 'Language code',
|
||||
},
|
||||
date: {
|
||||
type: 'string',
|
||||
description: 'Date in YYYY-MM-DD format for timestamp and daily_aggregation',
|
||||
},
|
||||
tz: {
|
||||
type: 'string',
|
||||
description: 'Timezone',
|
||||
},
|
||||
},
|
||||
required: ['action'],
|
||||
};
|
||||
|
||||
/**
|
||||
* Map user-friendly units to OpenWeather units.
|
||||
* Defaults to Celsius if not specified.
|
||||
|
|
@ -66,17 +110,11 @@ class OpenWeather extends Tool {
|
|||
'Units: "Celsius", "Kelvin", or "Fahrenheit" (default: Celsius). ' +
|
||||
'For timestamp action, use "date" in YYYY-MM-DD format.';
|
||||
|
||||
schema = z.object({
|
||||
action: z.enum(['help', 'current_forecast', 'timestamp', 'daily_aggregation', 'overview']),
|
||||
city: z.string().optional(),
|
||||
lat: z.number().optional(),
|
||||
lon: z.number().optional(),
|
||||
exclude: z.string().optional(),
|
||||
units: z.enum(['Celsius', 'Kelvin', 'Fahrenheit']).optional(),
|
||||
lang: z.string().optional(),
|
||||
date: z.string().optional(), // For timestamp and daily_aggregation
|
||||
tz: z.string().optional(),
|
||||
});
|
||||
schema = openWeatherJsonSchema;
|
||||
|
||||
static get jsonSchema() {
|
||||
return openWeatherJsonSchema;
|
||||
}
|
||||
|
||||
constructor(fields = {}) {
|
||||
super();
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
// Generates image using stable diffusion webui's api (automatic1111)
|
||||
const fs = require('fs');
|
||||
const { z } = require('zod');
|
||||
const path = require('path');
|
||||
const axios = require('axios');
|
||||
const sharp = require('sharp');
|
||||
|
|
@ -11,6 +10,23 @@ const { FileContext, ContentTypes } = require('librechat-data-provider');
|
|||
const { getBasePath } = require('@librechat/api');
|
||||
const paths = require('~/config/paths');
|
||||
|
||||
const stableDiffusionJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
prompt: {
|
||||
type: 'string',
|
||||
description:
|
||||
'Detailed keywords to describe the subject, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
},
|
||||
negative_prompt: {
|
||||
type: 'string',
|
||||
description:
|
||||
'Keywords we want to exclude from the final image, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
},
|
||||
},
|
||||
required: ['prompt', 'negative_prompt'],
|
||||
};
|
||||
|
||||
const displayMessage =
|
||||
"Stable Diffusion displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
|
|
@ -46,18 +62,11 @@ class StableDiffusionAPI extends Tool {
|
|||
// - Generate images only once per human query unless explicitly requested by the user`;
|
||||
this.description =
|
||||
"You can generate images using text with 'stable-diffusion'. This tool is exclusively for visual content.";
|
||||
this.schema = z.object({
|
||||
prompt: z
|
||||
.string()
|
||||
.describe(
|
||||
'Detailed keywords to describe the subject, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
),
|
||||
negative_prompt: z
|
||||
.string()
|
||||
.describe(
|
||||
'Keywords we want to exclude from the final image, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
),
|
||||
});
|
||||
this.schema = stableDiffusionJsonSchema;
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return stableDiffusionJsonSchema;
|
||||
}
|
||||
|
||||
replaceNewLinesWithSpaces(inputString) {
|
||||
|
|
|
|||
|
|
@ -1,8 +1,75 @@
|
|||
const { z } = require('zod');
|
||||
const { ProxyAgent, fetch } = require('undici');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
const tavilySearchJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
query: {
|
||||
type: 'string',
|
||||
minLength: 1,
|
||||
description: 'The search query string.',
|
||||
},
|
||||
max_results: {
|
||||
type: 'number',
|
||||
minimum: 1,
|
||||
maximum: 10,
|
||||
description: 'The maximum number of search results to return. Defaults to 5.',
|
||||
},
|
||||
search_depth: {
|
||||
type: 'string',
|
||||
enum: ['basic', 'advanced'],
|
||||
description:
|
||||
'The depth of the search, affecting result quality and response time (`basic` or `advanced`). Default is basic for quick results and advanced for indepth high quality results but longer response time. Advanced calls equals 2 requests.',
|
||||
},
|
||||
include_images: {
|
||||
type: 'boolean',
|
||||
description:
|
||||
'Whether to include a list of query-related images in the response. Default is False.',
|
||||
},
|
||||
include_answer: {
|
||||
type: 'boolean',
|
||||
description: 'Whether to include answers in the search results. Default is False.',
|
||||
},
|
||||
include_raw_content: {
|
||||
type: 'boolean',
|
||||
description: 'Whether to include raw content in the search results. Default is False.',
|
||||
},
|
||||
include_domains: {
|
||||
type: 'array',
|
||||
items: { type: 'string' },
|
||||
description: 'A list of domains to specifically include in the search results.',
|
||||
},
|
||||
exclude_domains: {
|
||||
type: 'array',
|
||||
items: { type: 'string' },
|
||||
description: 'A list of domains to specifically exclude from the search results.',
|
||||
},
|
||||
topic: {
|
||||
type: 'string',
|
||||
enum: ['general', 'news', 'finance'],
|
||||
description:
|
||||
'The category of the search. Use news ONLY if query SPECIFCALLY mentions the word "news".',
|
||||
},
|
||||
time_range: {
|
||||
type: 'string',
|
||||
enum: ['day', 'week', 'month', 'year', 'd', 'w', 'm', 'y'],
|
||||
description: 'The time range back from the current date to filter results.',
|
||||
},
|
||||
days: {
|
||||
type: 'number',
|
||||
minimum: 1,
|
||||
description: 'Number of days back from the current date to include. Only if topic is news.',
|
||||
},
|
||||
include_image_descriptions: {
|
||||
type: 'boolean',
|
||||
description:
|
||||
'When include_images is true, also add a descriptive text for each image. Default is false.',
|
||||
},
|
||||
},
|
||||
required: ['query'],
|
||||
};
|
||||
|
||||
class TavilySearchResults extends Tool {
|
||||
static lc_name() {
|
||||
return 'TavilySearchResults';
|
||||
|
|
@ -20,64 +87,11 @@ class TavilySearchResults extends Tool {
|
|||
this.description =
|
||||
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.';
|
||||
|
||||
this.schema = z.object({
|
||||
query: z.string().min(1).describe('The search query string.'),
|
||||
max_results: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(10)
|
||||
.optional()
|
||||
.describe('The maximum number of search results to return. Defaults to 5.'),
|
||||
search_depth: z
|
||||
.enum(['basic', 'advanced'])
|
||||
.optional()
|
||||
.describe(
|
||||
'The depth of the search, affecting result quality and response time (`basic` or `advanced`). Default is basic for quick results and advanced for indepth high quality results but longer response time. Advanced calls equals 2 requests.',
|
||||
),
|
||||
include_images: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe(
|
||||
'Whether to include a list of query-related images in the response. Default is False.',
|
||||
),
|
||||
include_answer: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe('Whether to include answers in the search results. Default is False.'),
|
||||
include_raw_content: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe('Whether to include raw content in the search results. Default is False.'),
|
||||
include_domains: z
|
||||
.array(z.string())
|
||||
.optional()
|
||||
.describe('A list of domains to specifically include in the search results.'),
|
||||
exclude_domains: z
|
||||
.array(z.string())
|
||||
.optional()
|
||||
.describe('A list of domains to specifically exclude from the search results.'),
|
||||
topic: z
|
||||
.enum(['general', 'news', 'finance'])
|
||||
.optional()
|
||||
.describe(
|
||||
'The category of the search. Use news ONLY if query SPECIFCALLY mentions the word "news".',
|
||||
),
|
||||
time_range: z
|
||||
.enum(['day', 'week', 'month', 'year', 'd', 'w', 'm', 'y'])
|
||||
.optional()
|
||||
.describe('The time range back from the current date to filter results.'),
|
||||
days: z
|
||||
.number()
|
||||
.min(1)
|
||||
.optional()
|
||||
.describe('Number of days back from the current date to include. Only if topic is news.'),
|
||||
include_image_descriptions: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe(
|
||||
'When include_images is true, also add a descriptive text for each image. Default is false.',
|
||||
),
|
||||
});
|
||||
this.schema = tavilySearchJsonSchema;
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return tavilySearchJsonSchema;
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
|
|
@ -89,12 +103,7 @@ class TavilySearchResults extends Tool {
|
|||
}
|
||||
|
||||
async _call(input) {
|
||||
const validationResult = this.schema.safeParse(input);
|
||||
if (!validationResult.success) {
|
||||
throw new Error(`Validation failed: ${JSON.stringify(validationResult.error.issues)}`);
|
||||
}
|
||||
|
||||
const { query, ...rest } = validationResult.data;
|
||||
const { query, ...rest } = input;
|
||||
|
||||
const requestBody = {
|
||||
api_key: this.apiKey,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,19 @@
|
|||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
const traversaalSearchJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
query: {
|
||||
type: 'string',
|
||||
description:
|
||||
"A properly written sentence to be interpreted by an AI to search the web according to the user's request.",
|
||||
},
|
||||
},
|
||||
required: ['query'],
|
||||
};
|
||||
|
||||
/**
|
||||
* Tool for the Traversaal AI search API, Ares.
|
||||
*/
|
||||
|
|
@ -17,17 +28,15 @@ class TraversaalSearch extends Tool {
|
|||
Useful for when you need to answer questions about current events. Input should be a search query.`;
|
||||
this.description_for_model =
|
||||
'\'Please create a specific sentence for the AI to understand and use as a query to search the web based on the user\'s request. For example, "Find information about the highest mountains in the world." or "Show me the latest news articles about climate change and its impact on polar ice caps."\'';
|
||||
this.schema = z.object({
|
||||
query: z
|
||||
.string()
|
||||
.describe(
|
||||
"A properly written sentence to be interpreted by an AI to search the web according to the user's request.",
|
||||
),
|
||||
});
|
||||
this.schema = traversaalSearchJsonSchema;
|
||||
|
||||
this.apiKey = fields?.TRAVERSAAL_API_KEY ?? this.getApiKey();
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return traversaalSearchJsonSchema;
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const apiKey = getEnvironmentVariable('TRAVERSAAL_API_KEY');
|
||||
if (!apiKey && this.override) {
|
||||
|
|
|
|||
|
|
@ -1,9 +1,19 @@
|
|||
/* eslint-disable no-useless-escape */
|
||||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
const wolframJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
input: {
|
||||
type: 'string',
|
||||
description: 'Natural language query to WolframAlpha following the guidelines',
|
||||
},
|
||||
},
|
||||
required: ['input'],
|
||||
};
|
||||
|
||||
class WolframAlphaAPI extends Tool {
|
||||
constructor(fields) {
|
||||
super();
|
||||
|
|
@ -41,9 +51,11 @@ class WolframAlphaAPI extends Tool {
|
|||
// -- Do not explain each step unless user input is needed. Proceed directly to making a better API call based on the available assumptions.`;
|
||||
this.description = `WolframAlpha offers computation, math, curated knowledge, and real-time data. It handles natural language queries and performs complex calculations.
|
||||
Follow the guidelines to get the best results.`;
|
||||
this.schema = z.object({
|
||||
input: z.string().describe('Natural language query to WolframAlpha following the guidelines'),
|
||||
});
|
||||
this.schema = wolframJsonSchema;
|
||||
}
|
||||
|
||||
static get jsonSchema() {
|
||||
return wolframJsonSchema;
|
||||
}
|
||||
|
||||
async fetchRawText(url) {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
|
@ -7,6 +6,18 @@ const { Tools, EToolResources } = require('librechat-data-provider');
|
|||
const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions');
|
||||
const { getFiles } = require('~/models');
|
||||
|
||||
const fileSearchJsonSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
query: {
|
||||
type: 'string',
|
||||
description:
|
||||
"A natural language query to search for relevant information in the files. Be specific and use keywords related to the information you're looking for. The query will be used for semantic similarity matching against the file contents.",
|
||||
},
|
||||
},
|
||||
required: ['query'],
|
||||
};
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {Object} options
|
||||
|
|
@ -182,15 +193,9 @@ Use the EXACT anchor markers shown below (copy them verbatim) immediately after
|
|||
**ALWAYS mention the filename in your text before the citation marker. NEVER use markdown links or footnotes.**`
|
||||
: ''
|
||||
}`,
|
||||
schema: z.object({
|
||||
query: z
|
||||
.string()
|
||||
.describe(
|
||||
"A natural language query to search for relevant information in the files. Be specific and use keywords related to the information you're looking for. The query will be used for semantic similarity matching against the file contents.",
|
||||
),
|
||||
}),
|
||||
schema: fileSearchJsonSchema,
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
module.exports = { createFileSearchTool, primeFiles };
|
||||
module.exports = { createFileSearchTool, primeFiles, fileSearchJsonSchema };
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ const {
|
|||
mcpToolPattern,
|
||||
loadWebSearchAuth,
|
||||
buildImageToolContext,
|
||||
buildWebSearchContext,
|
||||
} = require('@librechat/api');
|
||||
const { getMCPServersRegistry } = require('~/config');
|
||||
const {
|
||||
|
|
@ -19,7 +20,6 @@ const {
|
|||
Permissions,
|
||||
EToolResources,
|
||||
PermissionTypes,
|
||||
replaceSpecialVars,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
availableTools,
|
||||
|
|
@ -325,24 +325,7 @@ const loadTools = async ({
|
|||
});
|
||||
const { onSearchResults, onGetHighlights } = options?.[Tools.web_search] ?? {};
|
||||
requestedTools[tool] = async () => {
|
||||
toolContextMap[tool] = `# \`${tool}\`:
|
||||
Current Date & Time: ${replaceSpecialVars({ text: '{{iso_datetime}}' })}
|
||||
|
||||
**Execute immediately without preface.** After search, provide a brief summary addressing the query directly, then structure your response with clear Markdown formatting (## headers, lists, tables). Cite sources properly, tailor tone to query type, and provide comprehensive details.
|
||||
|
||||
**CITATION FORMAT - UNICODE ESCAPE SEQUENCES ONLY:**
|
||||
Use these EXACT escape sequences (copy verbatim): \\ue202 (before each anchor), \\ue200 (group start), \\ue201 (group end), \\ue203 (highlight start), \\ue204 (highlight end)
|
||||
|
||||
Anchor pattern: \\ue202turn{N}{type}{index} where N=turn number, type=search|news|image|ref, index=0,1,2...
|
||||
|
||||
**Examples (copy these exactly):**
|
||||
- Single: "Statement.\\ue202turn0search0"
|
||||
- Multiple: "Statement.\\ue202turn0search0\\ue202turn0news1"
|
||||
- Group: "Statement. \\ue200\\ue202turn0search0\\ue202turn0news1\\ue201"
|
||||
- Highlight: "\\ue203Cited text.\\ue204\\ue202turn0search0"
|
||||
- Image: "See photo\\ue202turn0image0."
|
||||
|
||||
**CRITICAL:** Output escape sequences EXACTLY as shown. Do NOT substitute with † or other symbols. Place anchors AFTER punctuation. Cite every non-obvious fact/quote. NEVER use markdown links, [1], footnotes, or HTML tags.`.trim();
|
||||
toolContextMap[tool] = buildWebSearchContext();
|
||||
return createSearchTool({
|
||||
...result.authResult,
|
||||
onSearchResults,
|
||||
|
|
|
|||
1
api/cache/banViolation.js
vendored
1
api/cache/banViolation.js
vendored
|
|
@ -55,6 +55,7 @@ const banViolation = async (req, res, errorMessage) => {
|
|||
|
||||
res.clearCookie('refreshToken');
|
||||
res.clearCookie('openid_access_token');
|
||||
res.clearCookie('openid_id_token');
|
||||
res.clearCookie('openid_user_id');
|
||||
res.clearCookie('token_provider');
|
||||
|
||||
|
|
|
|||
5
api/cache/getLogStores.js
vendored
5
api/cache/getLogStores.js
vendored
|
|
@ -37,6 +37,7 @@ const namespaces = {
|
|||
[CacheKeys.ROLES]: standardCache(CacheKeys.ROLES),
|
||||
[CacheKeys.APP_CONFIG]: standardCache(CacheKeys.APP_CONFIG),
|
||||
[CacheKeys.CONFIG_STORE]: standardCache(CacheKeys.CONFIG_STORE),
|
||||
[CacheKeys.TOOL_CACHE]: standardCache(CacheKeys.TOOL_CACHE),
|
||||
[CacheKeys.PENDING_REQ]: standardCache(CacheKeys.PENDING_REQ),
|
||||
[CacheKeys.ENCODED_DOMAINS]: new Keyv({ store: keyvMongo, namespace: CacheKeys.ENCODED_DOMAINS }),
|
||||
[CacheKeys.ABORT_KEYS]: standardCache(CacheKeys.ABORT_KEYS, Time.TEN_MINUTES),
|
||||
|
|
@ -51,6 +52,10 @@ const namespaces = {
|
|||
CacheKeys.OPENID_EXCHANGED_TOKENS,
|
||||
Time.TEN_MINUTES,
|
||||
),
|
||||
[CacheKeys.ADMIN_OAUTH_EXCHANGE]: standardCache(
|
||||
CacheKeys.ADMIN_OAUTH_EXCHANGE,
|
||||
Time.THIRTY_SECONDS,
|
||||
),
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -40,6 +40,10 @@ if (!cached) {
|
|||
cached = global.mongoose = { conn: null, promise: null };
|
||||
}
|
||||
|
||||
mongoose.connection.on('error', (err) => {
|
||||
logger.error('[connectDb] MongoDB connection error:', err);
|
||||
});
|
||||
|
||||
async function connectDb() {
|
||||
if (cached.conn && cached.conn?._readyState === 1) {
|
||||
return cached.conn;
|
||||
|
|
|
|||
|
|
@ -13,6 +13,11 @@ const searchEnabled = isEnabled(process.env.SEARCH);
|
|||
const indexingDisabled = isEnabled(process.env.MEILI_NO_SYNC);
|
||||
let currentTimeout = null;
|
||||
|
||||
const defaultSyncThreshold = 1000;
|
||||
const syncThreshold = process.env.MEILI_SYNC_THRESHOLD
|
||||
? parseInt(process.env.MEILI_SYNC_THRESHOLD, 10)
|
||||
: defaultSyncThreshold;
|
||||
|
||||
class MeiliSearchClient {
|
||||
static instance = null;
|
||||
|
||||
|
|
@ -221,25 +226,25 @@ async function performSync(flowManager, flowId, flowType) {
|
|||
}
|
||||
|
||||
// Check if we need to sync messages
|
||||
logger.info('[indexSync] Requesting message sync progress...');
|
||||
const messageProgress = await Message.getSyncProgress();
|
||||
if (!messageProgress.isComplete || settingsUpdated) {
|
||||
logger.info(
|
||||
`[indexSync] Messages need syncing: ${messageProgress.totalProcessed}/${messageProgress.totalDocuments} indexed`,
|
||||
);
|
||||
|
||||
// Check if we should do a full sync or incremental
|
||||
const messageCount = await Message.countDocuments();
|
||||
const messageCount = messageProgress.totalDocuments;
|
||||
const messagesIndexed = messageProgress.totalProcessed;
|
||||
const syncThreshold = parseInt(process.env.MEILI_SYNC_THRESHOLD || '1000', 10);
|
||||
const unindexedMessages = messageCount - messagesIndexed;
|
||||
|
||||
if (messageCount - messagesIndexed > syncThreshold) {
|
||||
logger.info('[indexSync] Starting full message sync due to large difference');
|
||||
await Message.syncWithMeili();
|
||||
messagesSync = true;
|
||||
} else if (messageCount !== messagesIndexed) {
|
||||
logger.warn('[indexSync] Messages out of sync, performing incremental sync');
|
||||
if (settingsUpdated || unindexedMessages > syncThreshold) {
|
||||
logger.info(`[indexSync] Starting message sync (${unindexedMessages} unindexed)`);
|
||||
await Message.syncWithMeili();
|
||||
messagesSync = true;
|
||||
} else if (unindexedMessages > 0) {
|
||||
logger.info(
|
||||
`[indexSync] ${unindexedMessages} messages unindexed (below threshold: ${syncThreshold}, skipping)`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
logger.info(
|
||||
|
|
@ -254,18 +259,18 @@ async function performSync(flowManager, flowId, flowType) {
|
|||
`[indexSync] Conversations need syncing: ${convoProgress.totalProcessed}/${convoProgress.totalDocuments} indexed`,
|
||||
);
|
||||
|
||||
const convoCount = await Conversation.countDocuments();
|
||||
const convoCount = convoProgress.totalDocuments;
|
||||
const convosIndexed = convoProgress.totalProcessed;
|
||||
const syncThreshold = parseInt(process.env.MEILI_SYNC_THRESHOLD || '1000', 10);
|
||||
|
||||
if (convoCount - convosIndexed > syncThreshold) {
|
||||
logger.info('[indexSync] Starting full conversation sync due to large difference');
|
||||
await Conversation.syncWithMeili();
|
||||
convosSync = true;
|
||||
} else if (convoCount !== convosIndexed) {
|
||||
logger.warn('[indexSync] Convos out of sync, performing incremental sync');
|
||||
const unindexedConvos = convoCount - convosIndexed;
|
||||
if (settingsUpdated || unindexedConvos > syncThreshold) {
|
||||
logger.info(`[indexSync] Starting convos sync (${unindexedConvos} unindexed)`);
|
||||
await Conversation.syncWithMeili();
|
||||
convosSync = true;
|
||||
} else if (unindexedConvos > 0) {
|
||||
logger.info(
|
||||
`[indexSync] ${unindexedConvos} convos unindexed (below threshold: ${syncThreshold}, skipping)`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
logger.info(
|
||||
|
|
|
|||
465
api/db/indexSync.spec.js
Normal file
465
api/db/indexSync.spec.js
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
/**
|
||||
* Unit tests for performSync() function in indexSync.js
|
||||
*
|
||||
* Tests use real mongoose with mocked model methods, only mocking external calls.
|
||||
*/
|
||||
|
||||
const mongoose = require('mongoose');
|
||||
|
||||
// Mock only external dependencies (not internal classes/models)
|
||||
const mockLogger = {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
};
|
||||
|
||||
const mockMeiliHealth = jest.fn();
|
||||
const mockMeiliIndex = jest.fn();
|
||||
const mockBatchResetMeiliFlags = jest.fn();
|
||||
const mockIsEnabled = jest.fn();
|
||||
const mockGetLogStores = jest.fn();
|
||||
|
||||
// Create mock models that will be reused
|
||||
const createMockModel = (collectionName) => ({
|
||||
collection: { name: collectionName },
|
||||
getSyncProgress: jest.fn(),
|
||||
syncWithMeili: jest.fn(),
|
||||
countDocuments: jest.fn(),
|
||||
});
|
||||
|
||||
const originalMessageModel = mongoose.models.Message;
|
||||
const originalConversationModel = mongoose.models.Conversation;
|
||||
|
||||
// Mock external modules
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('meilisearch', () => ({
|
||||
MeiliSearch: jest.fn(() => ({
|
||||
health: mockMeiliHealth,
|
||||
index: mockMeiliIndex,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock('./utils', () => ({
|
||||
batchResetMeiliFlags: mockBatchResetMeiliFlags,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: mockIsEnabled,
|
||||
FlowStateManager: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/cache', () => ({
|
||||
getLogStores: mockGetLogStores,
|
||||
}));
|
||||
|
||||
// Set environment before module load
|
||||
process.env.MEILI_HOST = 'http://localhost:7700';
|
||||
process.env.MEILI_MASTER_KEY = 'test-key';
|
||||
process.env.SEARCH = 'true';
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000'; // Set threshold before module loads
|
||||
|
||||
describe('performSync() - syncThreshold logic', () => {
|
||||
const ORIGINAL_ENV = process.env;
|
||||
let Message;
|
||||
let Conversation;
|
||||
|
||||
beforeAll(() => {
|
||||
Message = createMockModel('messages');
|
||||
Conversation = createMockModel('conversations');
|
||||
|
||||
mongoose.models.Message = Message;
|
||||
mongoose.models.Conversation = Conversation;
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset all mocks
|
||||
jest.clearAllMocks();
|
||||
// Reset modules to ensure fresh load of indexSync.js and its top-level consts (like syncThreshold)
|
||||
jest.resetModules();
|
||||
|
||||
// Set up environment
|
||||
process.env = { ...ORIGINAL_ENV };
|
||||
process.env.MEILI_HOST = 'http://localhost:7700';
|
||||
process.env.MEILI_MASTER_KEY = 'test-key';
|
||||
process.env.SEARCH = 'true';
|
||||
delete process.env.MEILI_NO_SYNC;
|
||||
|
||||
// Re-ensure models are available in mongoose after resetModules
|
||||
// We must require mongoose again to get the fresh instance that indexSync will use
|
||||
const mongoose = require('mongoose');
|
||||
mongoose.models.Message = Message;
|
||||
mongoose.models.Conversation = Conversation;
|
||||
|
||||
// Mock isEnabled
|
||||
mockIsEnabled.mockImplementation((val) => val === 'true' || val === true);
|
||||
|
||||
// Mock MeiliSearch client responses
|
||||
mockMeiliHealth.mockResolvedValue({ status: 'available' });
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: ['user'] }),
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
mockBatchResetMeiliFlags.mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = ORIGINAL_ENV;
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
mongoose.models.Message = originalMessageModel;
|
||||
mongoose.models.Conversation = originalConversationModel;
|
||||
});
|
||||
|
||||
test('triggers sync when unindexed messages exceed syncThreshold', async () => {
|
||||
// Arrange: Set threshold before module load
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Arrange: 1050 unindexed messages > 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1150, // 1050 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync triggered because 1050 > 1000
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1150 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (1050 unindexed)',
|
||||
);
|
||||
|
||||
// Assert: Conversation sync NOT triggered (already complete)
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('skips sync when unindexed messages are below syncThreshold', async () => {
|
||||
// Arrange: 50 unindexed messages < 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync NOT triggered because 50 < 1000
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/150 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] 50 messages unindexed (below threshold: 1000, skipping)',
|
||||
);
|
||||
|
||||
// Assert: Conversation sync NOT triggered (already complete)
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('respects syncThreshold at boundary (exactly at threshold)', async () => {
|
||||
// Arrange: 1000 unindexed messages = 1000 threshold (NOT greater than)
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1100, // 1000 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 0,
|
||||
totalDocuments: 0,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync NOT triggered because 1000 is NOT > 1000
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1100 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] 1000 messages unindexed (below threshold: 1000, skipping)',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers sync when unindexed is threshold + 1', async () => {
|
||||
// Arrange: 1001 unindexed messages > 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1101, // 1001 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 0,
|
||||
totalDocuments: 0,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync triggered because 1001 > 1000
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1101 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (1001 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('uses totalDocuments from convoProgress for conversation sync decisions', async () => {
|
||||
// Arrange: Messages complete, conversations need sync
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 1100, // 1050 unindexed > 1000 threshold
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls (the optimization)
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Only conversation sync triggered
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Conversations need syncing: 50/1100 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting convos sync (1050 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('skips sync when collections are fully synced', async () => {
|
||||
// Arrange: Everything already synced
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: No sync triggered
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Correct logs
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Messages are fully synced: 100/100');
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Conversations are fully synced: 50/50',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers message sync when settingsUpdated even if below syncThreshold', async () => {
|
||||
// Arrange: Only 50 unindexed messages (< 1000 threshold), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Message sync triggered despite being below threshold (50 < 1000)
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (50 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers conversation sync when settingsUpdated even if below syncThreshold', async () => {
|
||||
// Arrange: Messages complete, conversations have 50 unindexed (< 1000 threshold), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 100, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Conversation sync triggered despite being below threshold (50 < 1000)
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Starting convos sync (50 unindexed)');
|
||||
});
|
||||
|
||||
test('triggers both message and conversation sync when settingsUpdated even if both below syncThreshold', async () => {
|
||||
// Arrange: Set threshold before module load
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Arrange: Both have documents below threshold (50 each), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 100, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Both syncs triggered despite both being below threshold
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (50 unindexed)',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Starting convos sync (50 unindexed)');
|
||||
});
|
||||
});
|
||||
|
|
@ -26,7 +26,7 @@ async function batchResetMeiliFlags(collection) {
|
|||
try {
|
||||
while (hasMore) {
|
||||
const docs = await collection
|
||||
.find({ expiredAt: null, _meiliIndex: true }, { projection: { _id: 1 } })
|
||||
.find({ expiredAt: null, _meiliIndex: { $ne: false } }, { projection: { _id: 1 } })
|
||||
.limit(BATCH_SIZE)
|
||||
.toArray();
|
||||
|
||||
|
|
|
|||
|
|
@ -265,8 +265,8 @@ describe('batchResetMeiliFlags', () => {
|
|||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
// Only one document has _meiliIndex: true
|
||||
expect(result).toBe(1);
|
||||
// both documents should be updated
|
||||
expect(result).toBe(2);
|
||||
});
|
||||
|
||||
it('should handle mixed document states correctly', async () => {
|
||||
|
|
@ -275,16 +275,18 @@ describe('batchResetMeiliFlags', () => {
|
|||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: false },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: new Date(), _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: null },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(2);
|
||||
expect(result).toBe(4);
|
||||
|
||||
const flaggedDocs = await testCollection
|
||||
.find({ expiredAt: null, _meiliIndex: false })
|
||||
.toArray();
|
||||
expect(flaggedDocs).toHaveLength(3); // 2 were updated, 1 was already false
|
||||
expect(flaggedDocs).toHaveLength(5); // 4 were updated, 1 was already false
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -11,17 +11,15 @@ const {
|
|||
isEphemeralAgentId,
|
||||
encodeEphemeralAgentId,
|
||||
} = require('librechat-data-provider');
|
||||
const { GLOBAL_PROJECT_NAME, mcp_all, mcp_delimiter } =
|
||||
require('librechat-data-provider').Constants;
|
||||
const { mcp_all, mcp_delimiter } = require('librechat-data-provider').Constants;
|
||||
const {
|
||||
removeAgentFromAllProjects,
|
||||
removeAgentIdsFromProject,
|
||||
addAgentIdsToProject,
|
||||
getProjectByName,
|
||||
} = require('./Project');
|
||||
const { removeAllPermissions } = require('~/server/services/PermissionService');
|
||||
const { getMCPServerTools } = require('~/server/services/Config');
|
||||
const { Agent, AclEntry } = require('~/db/models');
|
||||
const { Agent, AclEntry, User } = require('~/db/models');
|
||||
const { getActions } = require('./Action');
|
||||
|
||||
/**
|
||||
|
|
@ -591,15 +589,29 @@ const deleteAgent = async (searchParameter) => {
|
|||
const agent = await Agent.findOneAndDelete(searchParameter);
|
||||
if (agent) {
|
||||
await removeAgentFromAllProjects(agent.id);
|
||||
await removeAllPermissions({
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
});
|
||||
await Promise.all([
|
||||
removeAllPermissions({
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
}),
|
||||
removeAllPermissions({
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
resourceId: agent._id,
|
||||
}),
|
||||
]);
|
||||
try {
|
||||
await Agent.updateMany({ 'edges.to': agent.id }, { $pull: { edges: { to: agent.id } } });
|
||||
} catch (error) {
|
||||
logger.error('[deleteAgent] Error removing agent from handoff edges', error);
|
||||
}
|
||||
try {
|
||||
await User.updateMany(
|
||||
{ 'favorites.agentId': agent.id },
|
||||
{ $pull: { favorites: { agentId: agent.id } } },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[deleteAgent] Error removing agent from user favorites', error);
|
||||
}
|
||||
}
|
||||
return agent;
|
||||
};
|
||||
|
|
@ -625,10 +637,19 @@ const deleteUserAgents = async (userId) => {
|
|||
}
|
||||
|
||||
await AclEntry.deleteMany({
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceType: { $in: [ResourceType.AGENT, ResourceType.REMOTE_AGENT] },
|
||||
resourceId: { $in: agentObjectIds },
|
||||
});
|
||||
|
||||
try {
|
||||
await User.updateMany(
|
||||
{ 'favorites.agentId': { $in: agentIds } },
|
||||
{ $pull: { favorites: { agentId: { $in: agentIds } } } },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[deleteUserAgents] Error removing agents from user favorites', error);
|
||||
}
|
||||
|
||||
await Agent.deleteMany({ author: userId });
|
||||
} catch (error) {
|
||||
logger.error('[deleteUserAgents] General error:', error);
|
||||
|
|
@ -735,59 +756,6 @@ const getListAgentsByAccess = async ({
|
|||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Get all agents.
|
||||
* @deprecated Use getListAgentsByAccess for ACL-aware agent listing
|
||||
* @param {Object} searchParameter - The search parameters to find matching agents.
|
||||
* @param {string} searchParameter.author - The user ID of the agent's author.
|
||||
* @returns {Promise<Object>} A promise that resolves to an object containing the agents data and pagination info.
|
||||
*/
|
||||
const getListAgents = async (searchParameter) => {
|
||||
const { author, ...otherParams } = searchParameter;
|
||||
|
||||
let query = Object.assign({ author }, otherParams);
|
||||
|
||||
const globalProject = await getProjectByName(GLOBAL_PROJECT_NAME, ['agentIds']);
|
||||
if (globalProject && (globalProject.agentIds?.length ?? 0) > 0) {
|
||||
const globalQuery = { id: { $in: globalProject.agentIds }, ...otherParams };
|
||||
delete globalQuery.author;
|
||||
query = { $or: [globalQuery, query] };
|
||||
}
|
||||
const agents = (
|
||||
await Agent.find(query, {
|
||||
id: 1,
|
||||
_id: 1,
|
||||
name: 1,
|
||||
avatar: 1,
|
||||
author: 1,
|
||||
projectIds: 1,
|
||||
description: 1,
|
||||
// @deprecated - isCollaborative replaced by ACL permissions
|
||||
isCollaborative: 1,
|
||||
category: 1,
|
||||
}).lean()
|
||||
).map((agent) => {
|
||||
if (agent.author?.toString() !== author) {
|
||||
delete agent.author;
|
||||
}
|
||||
if (agent.author) {
|
||||
agent.author = agent.author.toString();
|
||||
}
|
||||
return agent;
|
||||
});
|
||||
|
||||
const hasMore = agents.length > 0;
|
||||
const firstId = agents.length > 0 ? agents[0].id : null;
|
||||
const lastId = agents.length > 0 ? agents[agents.length - 1].id : null;
|
||||
|
||||
return {
|
||||
data: agents,
|
||||
has_more: hasMore,
|
||||
first_id: firstId,
|
||||
last_id: lastId,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates the projects associated with an agent, adding and removing project IDs as specified.
|
||||
* This function also updates the corresponding projects to include or exclude the agent ID.
|
||||
|
|
@ -953,12 +921,11 @@ module.exports = {
|
|||
updateAgent,
|
||||
deleteAgent,
|
||||
deleteUserAgents,
|
||||
getListAgents,
|
||||
revertAgentVersion,
|
||||
updateAgentProjects,
|
||||
countPromotedAgents,
|
||||
addAgentResourceFile,
|
||||
getListAgentsByAccess,
|
||||
removeAgentResourceFiles,
|
||||
generateActionMetadataHash,
|
||||
countPromotedAgents,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -22,17 +22,17 @@ const {
|
|||
createAgent,
|
||||
updateAgent,
|
||||
deleteAgent,
|
||||
getListAgents,
|
||||
getListAgentsByAccess,
|
||||
deleteUserAgents,
|
||||
revertAgentVersion,
|
||||
updateAgentProjects,
|
||||
addAgentResourceFile,
|
||||
getListAgentsByAccess,
|
||||
removeAgentResourceFiles,
|
||||
generateActionMetadataHash,
|
||||
} = require('./Agent');
|
||||
const permissionService = require('~/server/services/PermissionService');
|
||||
const { getCachedTools, getMCPServerTools } = require('~/server/services/Config');
|
||||
const { AclEntry } = require('~/db/models');
|
||||
const { AclEntry, User } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* @type {import('mongoose').Model<import('@librechat/data-schemas').IAgent>}
|
||||
|
|
@ -59,6 +59,7 @@ describe('models/Agent', () => {
|
|||
|
||||
beforeEach(async () => {
|
||||
await Agent.deleteMany({});
|
||||
await User.deleteMany({});
|
||||
});
|
||||
|
||||
test('should add tool_resource to tools if missing', async () => {
|
||||
|
|
@ -575,43 +576,488 @@ describe('models/Agent', () => {
|
|||
expect(sourceAgentAfter.edges).toHaveLength(0);
|
||||
});
|
||||
|
||||
test('should list agents by author', async () => {
|
||||
test('should remove agent from user favorites when agent is deleted', async () => {
|
||||
const agentId = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
||||
// Create agent
|
||||
await createAgent({
|
||||
id: agentId,
|
||||
name: 'Agent To Delete',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create user with the agent in favorites
|
||||
await User.create({
|
||||
_id: userId,
|
||||
name: 'Test User',
|
||||
email: `test-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agentId }, { model: 'gpt-4', endpoint: 'openAI' }],
|
||||
});
|
||||
|
||||
// Verify user has agent in favorites
|
||||
const userBefore = await User.findById(userId);
|
||||
expect(userBefore.favorites).toHaveLength(2);
|
||||
expect(userBefore.favorites.some((f) => f.agentId === agentId)).toBe(true);
|
||||
|
||||
// Delete the agent
|
||||
await deleteAgent({ id: agentId });
|
||||
|
||||
// Verify agent is deleted
|
||||
const agentAfterDelete = await getAgent({ id: agentId });
|
||||
expect(agentAfterDelete).toBeNull();
|
||||
|
||||
// Verify agent is removed from user favorites
|
||||
const userAfter = await User.findById(userId);
|
||||
expect(userAfter.favorites).toHaveLength(1);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agentId)).toBe(false);
|
||||
expect(userAfter.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
});
|
||||
|
||||
test('should remove agent from multiple users favorites when agent is deleted', async () => {
|
||||
const agentId = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const user1Id = new mongoose.Types.ObjectId();
|
||||
const user2Id = new mongoose.Types.ObjectId();
|
||||
|
||||
// Create agent
|
||||
await createAgent({
|
||||
id: agentId,
|
||||
name: 'Agent To Delete',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create two users with the agent in favorites
|
||||
await User.create({
|
||||
_id: user1Id,
|
||||
name: 'Test User 1',
|
||||
email: `test1-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agentId }],
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: user2Id,
|
||||
name: 'Test User 2',
|
||||
email: `test2-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agentId }, { agentId: `agent_${uuidv4()}` }],
|
||||
});
|
||||
|
||||
// Delete the agent
|
||||
await deleteAgent({ id: agentId });
|
||||
|
||||
// Verify agent is removed from both users' favorites
|
||||
const user1After = await User.findById(user1Id);
|
||||
const user2After = await User.findById(user2Id);
|
||||
|
||||
expect(user1After.favorites).toHaveLength(0);
|
||||
expect(user2After.favorites).toHaveLength(1);
|
||||
expect(user2After.favorites.some((f) => f.agentId === agentId)).toBe(false);
|
||||
});
|
||||
|
||||
test('should preserve other agents in database when one agent is deleted', async () => {
|
||||
const agentToDeleteId = `agent_${uuidv4()}`;
|
||||
const agentToKeep1Id = `agent_${uuidv4()}`;
|
||||
const agentToKeep2Id = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
|
||||
// Create multiple agents
|
||||
await createAgent({
|
||||
id: agentToDeleteId,
|
||||
name: 'Agent To Delete',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agentToKeep1Id,
|
||||
name: 'Agent To Keep 1',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agentToKeep2Id,
|
||||
name: 'Agent To Keep 2',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Verify all agents exist
|
||||
expect(await getAgent({ id: agentToDeleteId })).not.toBeNull();
|
||||
expect(await getAgent({ id: agentToKeep1Id })).not.toBeNull();
|
||||
expect(await getAgent({ id: agentToKeep2Id })).not.toBeNull();
|
||||
|
||||
// Delete one agent
|
||||
await deleteAgent({ id: agentToDeleteId });
|
||||
|
||||
// Verify only the deleted agent is removed, others remain intact
|
||||
expect(await getAgent({ id: agentToDeleteId })).toBeNull();
|
||||
const keptAgent1 = await getAgent({ id: agentToKeep1Id });
|
||||
const keptAgent2 = await getAgent({ id: agentToKeep2Id });
|
||||
expect(keptAgent1).not.toBeNull();
|
||||
expect(keptAgent1.name).toBe('Agent To Keep 1');
|
||||
expect(keptAgent2).not.toBeNull();
|
||||
expect(keptAgent2.name).toBe('Agent To Keep 2');
|
||||
});
|
||||
|
||||
test('should preserve other agents in user favorites when one agent is deleted', async () => {
|
||||
const agentToDeleteId = `agent_${uuidv4()}`;
|
||||
const agentToKeep1Id = `agent_${uuidv4()}`;
|
||||
const agentToKeep2Id = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
||||
// Create multiple agents
|
||||
await createAgent({
|
||||
id: agentToDeleteId,
|
||||
name: 'Agent To Delete',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agentToKeep1Id,
|
||||
name: 'Agent To Keep 1',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agentToKeep2Id,
|
||||
name: 'Agent To Keep 2',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create user with all three agents in favorites
|
||||
await User.create({
|
||||
_id: userId,
|
||||
name: 'Test User',
|
||||
email: `test-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [
|
||||
{ agentId: agentToDeleteId },
|
||||
{ agentId: agentToKeep1Id },
|
||||
{ agentId: agentToKeep2Id },
|
||||
],
|
||||
});
|
||||
|
||||
// Verify user has all three agents in favorites
|
||||
const userBefore = await User.findById(userId);
|
||||
expect(userBefore.favorites).toHaveLength(3);
|
||||
|
||||
// Delete one agent
|
||||
await deleteAgent({ id: agentToDeleteId });
|
||||
|
||||
// Verify only the deleted agent is removed from favorites
|
||||
const userAfter = await User.findById(userId);
|
||||
expect(userAfter.favorites).toHaveLength(2);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agentToDeleteId)).toBe(false);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agentToKeep1Id)).toBe(true);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agentToKeep2Id)).toBe(true);
|
||||
});
|
||||
|
||||
test('should not affect users who do not have deleted agent in favorites', async () => {
|
||||
const agentToDeleteId = `agent_${uuidv4()}`;
|
||||
const otherAgentId = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const userWithDeletedAgentId = new mongoose.Types.ObjectId();
|
||||
const userWithoutDeletedAgentId = new mongoose.Types.ObjectId();
|
||||
|
||||
// Create agents
|
||||
await createAgent({
|
||||
id: agentToDeleteId,
|
||||
name: 'Agent To Delete',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: otherAgentId,
|
||||
name: 'Other Agent',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create user with the agent to be deleted
|
||||
await User.create({
|
||||
_id: userWithDeletedAgentId,
|
||||
name: 'User With Deleted Agent',
|
||||
email: `user1-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agentToDeleteId }, { model: 'gpt-4', endpoint: 'openAI' }],
|
||||
});
|
||||
|
||||
// Create user without the agent to be deleted
|
||||
await User.create({
|
||||
_id: userWithoutDeletedAgentId,
|
||||
name: 'User Without Deleted Agent',
|
||||
email: `user2-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: otherAgentId }, { model: 'claude-3', endpoint: 'anthropic' }],
|
||||
});
|
||||
|
||||
// Delete the agent
|
||||
await deleteAgent({ id: agentToDeleteId });
|
||||
|
||||
// Verify user with deleted agent has it removed
|
||||
const userWithDeleted = await User.findById(userWithDeletedAgentId);
|
||||
expect(userWithDeleted.favorites).toHaveLength(1);
|
||||
expect(userWithDeleted.favorites.some((f) => f.agentId === agentToDeleteId)).toBe(false);
|
||||
expect(userWithDeleted.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
|
||||
// Verify user without deleted agent is completely unaffected
|
||||
const userWithoutDeleted = await User.findById(userWithoutDeletedAgentId);
|
||||
expect(userWithoutDeleted.favorites).toHaveLength(2);
|
||||
expect(userWithoutDeleted.favorites.some((f) => f.agentId === otherAgentId)).toBe(true);
|
||||
expect(userWithoutDeleted.favorites.some((f) => f.model === 'claude-3')).toBe(true);
|
||||
});
|
||||
|
||||
test('should remove all user agents from favorites when deleteUserAgents is called', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const otherAuthorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
||||
const agentIds = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
const id = `agent_${uuidv4()}`;
|
||||
agentIds.push(id);
|
||||
await createAgent({
|
||||
id,
|
||||
name: `Agent ${i}`,
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
}
|
||||
const agent1Id = `agent_${uuidv4()}`;
|
||||
const agent2Id = `agent_${uuidv4()}`;
|
||||
const otherAuthorAgentId = `agent_${uuidv4()}`;
|
||||
|
||||
for (let i = 0; i < 3; i++) {
|
||||
await createAgent({
|
||||
id: `other_agent_${uuidv4()}`,
|
||||
name: `Other Agent ${i}`,
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: otherAuthorId,
|
||||
});
|
||||
}
|
||||
// Create agents by the author to be deleted
|
||||
await createAgent({
|
||||
id: agent1Id,
|
||||
name: 'Author Agent 1',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
const result = await getListAgents({ author: authorId.toString() });
|
||||
await createAgent({
|
||||
id: agent2Id,
|
||||
name: 'Author Agent 2',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result.data).toBeDefined();
|
||||
expect(result.data).toHaveLength(5);
|
||||
expect(result.has_more).toBe(true);
|
||||
// Create agent by different author (should not be deleted)
|
||||
await createAgent({
|
||||
id: otherAuthorAgentId,
|
||||
name: 'Other Author Agent',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: otherAuthorId,
|
||||
});
|
||||
|
||||
for (const agent of result.data) {
|
||||
expect(agent.author).toBe(authorId.toString());
|
||||
}
|
||||
// Create user with all agents in favorites
|
||||
await User.create({
|
||||
_id: userId,
|
||||
name: 'Test User',
|
||||
email: `test-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [
|
||||
{ agentId: agent1Id },
|
||||
{ agentId: agent2Id },
|
||||
{ agentId: otherAuthorAgentId },
|
||||
{ model: 'gpt-4', endpoint: 'openAI' },
|
||||
],
|
||||
});
|
||||
|
||||
// Verify user has all favorites
|
||||
const userBefore = await User.findById(userId);
|
||||
expect(userBefore.favorites).toHaveLength(4);
|
||||
|
||||
// Delete all agents by the author
|
||||
await deleteUserAgents(authorId.toString());
|
||||
|
||||
// Verify author's agents are deleted from database
|
||||
expect(await getAgent({ id: agent1Id })).toBeNull();
|
||||
expect(await getAgent({ id: agent2Id })).toBeNull();
|
||||
|
||||
// Verify other author's agent still exists
|
||||
expect(await getAgent({ id: otherAuthorAgentId })).not.toBeNull();
|
||||
|
||||
// Verify user favorites: author's agents removed, others remain
|
||||
const userAfter = await User.findById(userId);
|
||||
expect(userAfter.favorites).toHaveLength(2);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agent1Id)).toBe(false);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === agent2Id)).toBe(false);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === otherAuthorAgentId)).toBe(true);
|
||||
expect(userAfter.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
});
|
||||
|
||||
test('should handle deleteUserAgents when agents are in multiple users favorites', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const user1Id = new mongoose.Types.ObjectId();
|
||||
const user2Id = new mongoose.Types.ObjectId();
|
||||
const user3Id = new mongoose.Types.ObjectId();
|
||||
|
||||
const agent1Id = `agent_${uuidv4()}`;
|
||||
const agent2Id = `agent_${uuidv4()}`;
|
||||
const unrelatedAgentId = `agent_${uuidv4()}`;
|
||||
|
||||
// Create agents by the author
|
||||
await createAgent({
|
||||
id: agent1Id,
|
||||
name: 'Author Agent 1',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agent2Id,
|
||||
name: 'Author Agent 2',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create users with various favorites configurations
|
||||
await User.create({
|
||||
_id: user1Id,
|
||||
name: 'User 1',
|
||||
email: `user1-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agent1Id }, { agentId: agent2Id }],
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: user2Id,
|
||||
name: 'User 2',
|
||||
email: `user2-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: agent1Id }, { model: 'claude-3', endpoint: 'anthropic' }],
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: user3Id,
|
||||
name: 'User 3',
|
||||
email: `user3-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: unrelatedAgentId }, { model: 'gpt-4', endpoint: 'openAI' }],
|
||||
});
|
||||
|
||||
// Delete all agents by the author
|
||||
await deleteUserAgents(authorId.toString());
|
||||
|
||||
// Verify all users' favorites are correctly updated
|
||||
const user1After = await User.findById(user1Id);
|
||||
expect(user1After.favorites).toHaveLength(0);
|
||||
|
||||
const user2After = await User.findById(user2Id);
|
||||
expect(user2After.favorites).toHaveLength(1);
|
||||
expect(user2After.favorites.some((f) => f.agentId === agent1Id)).toBe(false);
|
||||
expect(user2After.favorites.some((f) => f.model === 'claude-3')).toBe(true);
|
||||
|
||||
// User 3 should be completely unaffected
|
||||
const user3After = await User.findById(user3Id);
|
||||
expect(user3After.favorites).toHaveLength(2);
|
||||
expect(user3After.favorites.some((f) => f.agentId === unrelatedAgentId)).toBe(true);
|
||||
expect(user3After.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
});
|
||||
|
||||
test('should handle deleteUserAgents when user has no agents', async () => {
|
||||
const authorWithNoAgentsId = new mongoose.Types.ObjectId();
|
||||
const otherAuthorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
||||
const existingAgentId = `agent_${uuidv4()}`;
|
||||
|
||||
// Create agent by different author
|
||||
await createAgent({
|
||||
id: existingAgentId,
|
||||
name: 'Existing Agent',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: otherAuthorId,
|
||||
});
|
||||
|
||||
// Create user with favorites
|
||||
await User.create({
|
||||
_id: userId,
|
||||
name: 'Test User',
|
||||
email: `test-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ agentId: existingAgentId }, { model: 'gpt-4', endpoint: 'openAI' }],
|
||||
});
|
||||
|
||||
// Delete agents for user with no agents (should be a no-op)
|
||||
await deleteUserAgents(authorWithNoAgentsId.toString());
|
||||
|
||||
// Verify existing agent still exists
|
||||
expect(await getAgent({ id: existingAgentId })).not.toBeNull();
|
||||
|
||||
// Verify user favorites are unchanged
|
||||
const userAfter = await User.findById(userId);
|
||||
expect(userAfter.favorites).toHaveLength(2);
|
||||
expect(userAfter.favorites.some((f) => f.agentId === existingAgentId)).toBe(true);
|
||||
expect(userAfter.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
});
|
||||
|
||||
test('should handle deleteUserAgents when agents are not in any favorites', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
||||
const agent1Id = `agent_${uuidv4()}`;
|
||||
const agent2Id = `agent_${uuidv4()}`;
|
||||
|
||||
// Create agents by the author
|
||||
await createAgent({
|
||||
id: agent1Id,
|
||||
name: 'Agent 1',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
await createAgent({
|
||||
id: agent2Id,
|
||||
name: 'Agent 2',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
});
|
||||
|
||||
// Create user with favorites that don't include these agents
|
||||
await User.create({
|
||||
_id: userId,
|
||||
name: 'Test User',
|
||||
email: `test-${uuidv4()}@example.com`,
|
||||
provider: 'local',
|
||||
favorites: [{ model: 'gpt-4', endpoint: 'openAI' }],
|
||||
});
|
||||
|
||||
// Verify agents exist
|
||||
expect(await getAgent({ id: agent1Id })).not.toBeNull();
|
||||
expect(await getAgent({ id: agent2Id })).not.toBeNull();
|
||||
|
||||
// Delete all agents by the author
|
||||
await deleteUserAgents(authorId.toString());
|
||||
|
||||
// Verify agents are deleted
|
||||
expect(await getAgent({ id: agent1Id })).toBeNull();
|
||||
expect(await getAgent({ id: agent2Id })).toBeNull();
|
||||
|
||||
// Verify user favorites are unchanged
|
||||
const userAfter = await User.findById(userId);
|
||||
expect(userAfter.favorites).toHaveLength(1);
|
||||
expect(userAfter.favorites.some((f) => f.model === 'gpt-4')).toBe(true);
|
||||
});
|
||||
|
||||
test('should update agent projects', async () => {
|
||||
|
|
@ -733,26 +1179,6 @@ describe('models/Agent', () => {
|
|||
expect(result).toBe(expected);
|
||||
});
|
||||
|
||||
test('should handle getListAgents with invalid author format', async () => {
|
||||
try {
|
||||
const result = await getListAgents({ author: 'invalid-object-id' });
|
||||
expect(result.data).toEqual([]);
|
||||
} catch (error) {
|
||||
expect(error).toBeDefined();
|
||||
}
|
||||
});
|
||||
|
||||
test('should handle getListAgents with no agents', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const result = await getListAgents({ author: authorId.toString() });
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result.data).toEqual([]);
|
||||
expect(result.has_more).toBe(false);
|
||||
expect(result.first_id).toBeNull();
|
||||
expect(result.last_id).toBeNull();
|
||||
});
|
||||
|
||||
test('should handle updateAgentProjects with non-existent agent', async () => {
|
||||
const nonExistentId = `agent_${uuidv4()}`;
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
|
|
@ -2366,17 +2792,6 @@ describe('models/Agent', () => {
|
|||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
test('should handle getListAgents with no agents', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const result = await getListAgents({ author: authorId.toString() });
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result.data).toEqual([]);
|
||||
expect(result.has_more).toBe(false);
|
||||
expect(result.first_id).toBeNull();
|
||||
expect(result.last_id).toBeNull();
|
||||
});
|
||||
|
||||
test('should handle updateAgent with MongoDB operators mixed with direct updates', async () => {
|
||||
const agentId = `agent_${uuidv4()}`;
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
|
|
|
|||
|
|
@ -124,10 +124,15 @@ module.exports = {
|
|||
updateOperation,
|
||||
{
|
||||
new: true,
|
||||
upsert: true,
|
||||
upsert: metadata?.noUpsert !== true,
|
||||
},
|
||||
);
|
||||
|
||||
if (!conversation) {
|
||||
logger.debug('[saveConvo] Conversation not found, skipping update');
|
||||
return null;
|
||||
}
|
||||
|
||||
return conversation.toObject();
|
||||
} catch (error) {
|
||||
logger.error('[saveConvo] Error saving conversation', error);
|
||||
|
|
|
|||
|
|
@ -106,6 +106,47 @@ describe('Conversation Operations', () => {
|
|||
expect(result.conversationId).toBe(newConversationId);
|
||||
});
|
||||
|
||||
it('should not create a conversation when noUpsert is true and conversation does not exist', async () => {
|
||||
const nonExistentId = uuidv4();
|
||||
const result = await saveConvo(
|
||||
mockReq,
|
||||
{ conversationId: nonExistentId, title: 'Ghost Title' },
|
||||
{ noUpsert: true },
|
||||
);
|
||||
|
||||
expect(result).toBeNull();
|
||||
|
||||
const dbConvo = await Conversation.findOne({ conversationId: nonExistentId });
|
||||
expect(dbConvo).toBeNull();
|
||||
});
|
||||
|
||||
it('should update an existing conversation when noUpsert is true', async () => {
|
||||
await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
const result = await saveConvo(
|
||||
mockReq,
|
||||
{ conversationId: mockConversationData.conversationId, title: 'Updated Title' },
|
||||
{ noUpsert: true },
|
||||
);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result.title).toBe('Updated Title');
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
});
|
||||
|
||||
it('should still upsert by default when noUpsert is not provided', async () => {
|
||||
const newId = uuidv4();
|
||||
const result = await saveConvo(mockReq, {
|
||||
conversationId: newId,
|
||||
title: 'New Conversation',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result.conversationId).toBe(newId);
|
||||
expect(result.title).toBe('New Conversation');
|
||||
});
|
||||
|
||||
it('should handle unsetFields metadata', async () => {
|
||||
const metadata = {
|
||||
unsetFields: { someField: 1 },
|
||||
|
|
@ -122,7 +163,6 @@ describe('Conversation Operations', () => {
|
|||
|
||||
describe('isTemporary conversation handling', () => {
|
||||
it('should save a conversation with expiredAt when isTemporary is true', async () => {
|
||||
// Mock app config with 24 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
|
@ -135,7 +175,6 @@ describe('Conversation Operations', () => {
|
|||
expect(result.expiredAt).toBeDefined();
|
||||
expect(result.expiredAt).toBeInstanceOf(Date);
|
||||
|
||||
// Verify expiredAt is approximately 24 hours in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 24 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
|
|
@ -157,7 +196,6 @@ describe('Conversation Operations', () => {
|
|||
});
|
||||
|
||||
it('should save a conversation without expiredAt when isTemporary is not provided', async () => {
|
||||
// No isTemporary in body
|
||||
mockReq.body = {};
|
||||
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
|
@ -167,7 +205,6 @@ describe('Conversation Operations', () => {
|
|||
});
|
||||
|
||||
it('should use custom retention period from config', async () => {
|
||||
// Mock app config with 48 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 48;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
|
|
|||
|
|
@ -26,7 +26,8 @@ const getFiles = async (filter, _sortOptions, selectFields = { text: 0 }) => {
|
|||
};
|
||||
|
||||
/**
|
||||
* Retrieves tool files (files that are embedded or have a fileIdentifier) from an array of file IDs
|
||||
* Retrieves tool files (files that are embedded or have a fileIdentifier) from an array of file IDs.
|
||||
* Note: execute_code files are handled separately by getCodeGeneratedFiles.
|
||||
* @param {string[]} fileIds - Array of file_id strings to search for
|
||||
* @param {Set<EToolResources>} toolResourceSet - Optional filter for tool resources
|
||||
* @returns {Promise<Array<MongoFile>>} Files that match the criteria
|
||||
|
|
@ -37,21 +38,25 @@ const getToolFilesByIds = async (fileIds, toolResourceSet) => {
|
|||
}
|
||||
|
||||
try {
|
||||
const filter = {
|
||||
file_id: { $in: fileIds },
|
||||
$or: [],
|
||||
};
|
||||
const orConditions = [];
|
||||
|
||||
if (toolResourceSet.has(EToolResources.context)) {
|
||||
filter.$or.push({ text: { $exists: true, $ne: null }, context: FileContext.agents });
|
||||
orConditions.push({ text: { $exists: true, $ne: null }, context: FileContext.agents });
|
||||
}
|
||||
if (toolResourceSet.has(EToolResources.file_search)) {
|
||||
filter.$or.push({ embedded: true });
|
||||
orConditions.push({ embedded: true });
|
||||
}
|
||||
if (toolResourceSet.has(EToolResources.execute_code)) {
|
||||
filter.$or.push({ 'metadata.fileIdentifier': { $exists: true } });
|
||||
|
||||
if (orConditions.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const filter = {
|
||||
file_id: { $in: fileIds },
|
||||
context: { $ne: FileContext.execute_code }, // Exclude code-generated files
|
||||
$or: orConditions,
|
||||
};
|
||||
|
||||
const selectFields = { text: 0 };
|
||||
const sortOptions = { updatedAt: -1 };
|
||||
|
||||
|
|
@ -62,6 +67,70 @@ const getToolFilesByIds = async (fileIds, toolResourceSet) => {
|
|||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves files generated by code execution for a given conversation.
|
||||
* These files are stored locally with fileIdentifier metadata for code env re-upload.
|
||||
* @param {string} conversationId - The conversation ID to search for
|
||||
* @param {string[]} [messageIds] - Optional array of messageIds to filter by (for linear thread filtering)
|
||||
* @returns {Promise<Array<MongoFile>>} Files generated by code execution in the conversation
|
||||
*/
|
||||
const getCodeGeneratedFiles = async (conversationId, messageIds) => {
|
||||
if (!conversationId) {
|
||||
return [];
|
||||
}
|
||||
|
||||
/** messageIds are required for proper thread filtering of code-generated files */
|
||||
if (!messageIds || messageIds.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
try {
|
||||
const filter = {
|
||||
conversationId,
|
||||
context: FileContext.execute_code,
|
||||
messageId: { $exists: true, $in: messageIds },
|
||||
'metadata.fileIdentifier': { $exists: true },
|
||||
};
|
||||
|
||||
const selectFields = { text: 0 };
|
||||
const sortOptions = { createdAt: 1 };
|
||||
|
||||
return await getFiles(filter, sortOptions, selectFields);
|
||||
} catch (error) {
|
||||
logger.error('[getCodeGeneratedFiles] Error retrieving code generated files:', error);
|
||||
return [];
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves user-uploaded execute_code files (not code-generated) by their file IDs.
|
||||
* These are files with fileIdentifier metadata but context is NOT execute_code (e.g., agents or message_attachment).
|
||||
* File IDs should be collected from message.files arrays in the current thread.
|
||||
* @param {string[]} fileIds - Array of file IDs to fetch (from message.files in the thread)
|
||||
* @returns {Promise<Array<MongoFile>>} User-uploaded execute_code files
|
||||
*/
|
||||
const getUserCodeFiles = async (fileIds) => {
|
||||
if (!fileIds || fileIds.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
try {
|
||||
const filter = {
|
||||
file_id: { $in: fileIds },
|
||||
context: { $ne: FileContext.execute_code },
|
||||
'metadata.fileIdentifier': { $exists: true },
|
||||
};
|
||||
|
||||
const selectFields = { text: 0 };
|
||||
const sortOptions = { createdAt: 1 };
|
||||
|
||||
return await getFiles(filter, sortOptions, selectFields);
|
||||
} catch (error) {
|
||||
logger.error('[getUserCodeFiles] Error retrieving user code files:', error);
|
||||
return [];
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Creates a new file with a TTL of 1 hour.
|
||||
* @param {MongoFile} data - The file data to be created, must contain file_id.
|
||||
|
|
@ -169,6 +238,8 @@ module.exports = {
|
|||
findFileById,
|
||||
getFiles,
|
||||
getToolFilesByIds,
|
||||
getCodeGeneratedFiles,
|
||||
getUserCodeFiles,
|
||||
createFile,
|
||||
updateFile,
|
||||
updateFileUsage,
|
||||
|
|
|
|||
|
|
@ -138,11 +138,10 @@ const updateBalance = async ({ user, incrementValue, setValues }) => {
|
|||
|
||||
/** Method to calculate and set the tokenValue for a transaction */
|
||||
function calculateTokenValue(txn) {
|
||||
if (!txn.valueKey || !txn.tokenType) {
|
||||
txn.tokenValue = txn.rawAmount;
|
||||
}
|
||||
const { valueKey, tokenType, model, endpointTokenConfig } = txn;
|
||||
const multiplier = Math.abs(getMultiplier({ valueKey, tokenType, model, endpointTokenConfig }));
|
||||
const { valueKey, tokenType, model, endpointTokenConfig, inputTokenCount } = txn;
|
||||
const multiplier = Math.abs(
|
||||
getMultiplier({ valueKey, tokenType, model, endpointTokenConfig, inputTokenCount }),
|
||||
);
|
||||
txn.rate = multiplier;
|
||||
txn.tokenValue = txn.rawAmount * multiplier;
|
||||
if (txn.context && txn.tokenType === 'completion' && txn.context === 'incomplete') {
|
||||
|
|
@ -166,6 +165,7 @@ async function createAutoRefillTransaction(txData) {
|
|||
}
|
||||
const transaction = new Transaction(txData);
|
||||
transaction.endpointTokenConfig = txData.endpointTokenConfig;
|
||||
transaction.inputTokenCount = txData.inputTokenCount;
|
||||
calculateTokenValue(transaction);
|
||||
await transaction.save();
|
||||
|
||||
|
|
@ -200,6 +200,7 @@ async function createTransaction(_txData) {
|
|||
|
||||
const transaction = new Transaction(txData);
|
||||
transaction.endpointTokenConfig = txData.endpointTokenConfig;
|
||||
transaction.inputTokenCount = txData.inputTokenCount;
|
||||
calculateTokenValue(transaction);
|
||||
|
||||
await transaction.save();
|
||||
|
|
@ -231,10 +232,9 @@ async function createStructuredTransaction(_txData) {
|
|||
return;
|
||||
}
|
||||
|
||||
const transaction = new Transaction({
|
||||
...txData,
|
||||
endpointTokenConfig: txData.endpointTokenConfig,
|
||||
});
|
||||
const transaction = new Transaction(txData);
|
||||
transaction.endpointTokenConfig = txData.endpointTokenConfig;
|
||||
transaction.inputTokenCount = txData.inputTokenCount;
|
||||
|
||||
calculateStructuredTokenValue(transaction);
|
||||
|
||||
|
|
@ -266,10 +266,15 @@ function calculateStructuredTokenValue(txn) {
|
|||
return;
|
||||
}
|
||||
|
||||
const { model, endpointTokenConfig } = txn;
|
||||
const { model, endpointTokenConfig, inputTokenCount } = txn;
|
||||
|
||||
if (txn.tokenType === 'prompt') {
|
||||
const inputMultiplier = getMultiplier({ tokenType: 'prompt', model, endpointTokenConfig });
|
||||
const inputMultiplier = getMultiplier({
|
||||
tokenType: 'prompt',
|
||||
model,
|
||||
endpointTokenConfig,
|
||||
inputTokenCount,
|
||||
});
|
||||
const writeMultiplier =
|
||||
getCacheMultiplier({ cacheType: 'write', model, endpointTokenConfig }) ?? inputMultiplier;
|
||||
const readMultiplier =
|
||||
|
|
@ -304,7 +309,12 @@ function calculateStructuredTokenValue(txn) {
|
|||
|
||||
txn.rawAmount = -totalPromptTokens;
|
||||
} else if (txn.tokenType === 'completion') {
|
||||
const multiplier = getMultiplier({ tokenType: txn.tokenType, model, endpointTokenConfig });
|
||||
const multiplier = getMultiplier({
|
||||
tokenType: txn.tokenType,
|
||||
model,
|
||||
endpointTokenConfig,
|
||||
inputTokenCount,
|
||||
});
|
||||
txn.rate = Math.abs(multiplier);
|
||||
txn.tokenValue = -Math.abs(txn.rawAmount) * multiplier;
|
||||
txn.rawAmount = -Math.abs(txn.rawAmount);
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { getMultiplier, getCacheMultiplier, premiumTokenValues, tokenValues } = require('./tx');
|
||||
const { createTransaction, createStructuredTransaction } = require('./Transaction');
|
||||
const { Balance, Transaction } = require('~/db/models');
|
||||
|
||||
|
|
@ -564,3 +564,291 @@ describe('Transactions Config Tests', () => {
|
|||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
|
||||
describe('calculateTokenValue Edge Cases', () => {
|
||||
test('should derive multiplier from model when valueKey is not provided', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-4';
|
||||
const promptTokens = 1000;
|
||||
|
||||
const result = await createTransaction({
|
||||
user: userId,
|
||||
conversationId: 'test-no-valuekey',
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -promptTokens,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
});
|
||||
|
||||
const expectedRate = getMultiplier({ model, tokenType: 'prompt' });
|
||||
expect(result.rate).toBe(expectedRate);
|
||||
|
||||
const tx = await Transaction.findOne({ user: userId });
|
||||
expect(tx.tokenValue).toBe(-promptTokens * expectedRate);
|
||||
expect(tx.rate).toBe(expectedRate);
|
||||
});
|
||||
|
||||
test('should derive valueKey and apply correct rate for an unknown model with tokenType', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
await createTransaction({
|
||||
user: userId,
|
||||
conversationId: 'test-unknown-model',
|
||||
model: 'some-unrecognized-model-xyz',
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -500,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
});
|
||||
|
||||
const tx = await Transaction.findOne({ user: userId });
|
||||
expect(tx.rate).toBeDefined();
|
||||
expect(tx.rate).toBeGreaterThan(0);
|
||||
expect(tx.tokenValue).toBe(tx.rawAmount * tx.rate);
|
||||
});
|
||||
|
||||
test('should correctly apply model-derived multiplier without valueKey for completion', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const completionTokens = 500;
|
||||
|
||||
const result = await createTransaction({
|
||||
user: userId,
|
||||
conversationId: 'test-completion-no-valuekey',
|
||||
model,
|
||||
tokenType: 'completion',
|
||||
rawAmount: -completionTokens,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
});
|
||||
|
||||
const expectedRate = getMultiplier({ model, tokenType: 'completion' });
|
||||
expect(expectedRate).toBe(tokenValues[model].completion);
|
||||
expect(result.rate).toBe(expectedRate);
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(
|
||||
initialBalance - completionTokens * expectedRate,
|
||||
0,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Premium Token Pricing Integration Tests', () => {
|
||||
test('spendTokens should apply standard pricing when prompt tokens are below premium threshold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 100000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-premium-below',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const standardPromptRate = tokenValues[model].prompt;
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
const expectedCost =
|
||||
promptTokens * standardPromptRate + completionTokens * standardCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should apply premium pricing when prompt tokens exceed premium threshold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 250000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-premium-above',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
const expectedCost =
|
||||
promptTokens * premiumPromptRate + completionTokens * premiumCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should apply standard pricing at exactly the premium threshold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = premiumTokenValues[model].threshold;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-premium-exact',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const standardPromptRate = tokenValues[model].prompt;
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
const expectedCost =
|
||||
promptTokens * standardPromptRate + completionTokens * standardCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendStructuredTokens should apply premium pricing when total input tokens exceed threshold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-structured-premium',
|
||||
model,
|
||||
context: 'message',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 200000,
|
||||
write: 10000,
|
||||
read: 5000,
|
||||
},
|
||||
completionTokens: 1000,
|
||||
};
|
||||
|
||||
const totalInput =
|
||||
tokenUsage.promptTokens.input + tokenUsage.promptTokens.write + tokenUsage.promptTokens.read;
|
||||
|
||||
await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * premiumPromptRate +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * premiumCompletionRate;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(totalInput).toBeGreaterThan(premiumTokenValues[model].threshold);
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedTotalCost, 0);
|
||||
});
|
||||
|
||||
test('spendStructuredTokens should apply standard pricing when total input tokens are below threshold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-structured-standard',
|
||||
model,
|
||||
context: 'message',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 50000,
|
||||
write: 10000,
|
||||
read: 5000,
|
||||
},
|
||||
completionTokens: 1000,
|
||||
};
|
||||
|
||||
const totalInput =
|
||||
tokenUsage.promptTokens.input + tokenUsage.promptTokens.write + tokenUsage.promptTokens.read;
|
||||
|
||||
await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
const standardPromptRate = tokenValues[model].prompt;
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * standardPromptRate +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * standardCompletionRate;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(totalInput).toBeLessThanOrEqual(premiumTokenValues[model].threshold);
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedTotalCost, 0);
|
||||
});
|
||||
|
||||
test('non-premium models should not be affected by inputTokenCount regardless of prompt size', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-5';
|
||||
const promptTokens = 300000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-no-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const standardPromptRate = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const standardCompletionRate = getMultiplier({ model, tokenType: 'completion' });
|
||||
const expectedCost =
|
||||
promptTokens * standardPromptRate + completionTokens * standardCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -24,12 +24,14 @@ const spendTokens = async (txData, tokenUsage) => {
|
|||
},
|
||||
);
|
||||
let prompt, completion;
|
||||
const normalizedPromptTokens = Math.max(promptTokens ?? 0, 0);
|
||||
try {
|
||||
if (promptTokens !== undefined) {
|
||||
prompt = await createTransaction({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
rawAmount: promptTokens === 0 ? 0 : -Math.max(promptTokens, 0),
|
||||
rawAmount: promptTokens === 0 ? 0 : -normalizedPromptTokens,
|
||||
inputTokenCount: normalizedPromptTokens,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -38,6 +40,7 @@ const spendTokens = async (txData, tokenUsage) => {
|
|||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: completionTokens === 0 ? 0 : -Math.max(completionTokens, 0),
|
||||
inputTokenCount: normalizedPromptTokens,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -87,21 +90,31 @@ const spendStructuredTokens = async (txData, tokenUsage) => {
|
|||
let prompt, completion;
|
||||
try {
|
||||
if (promptTokens) {
|
||||
const { input = 0, write = 0, read = 0 } = promptTokens;
|
||||
const input = Math.max(promptTokens.input ?? 0, 0);
|
||||
const write = Math.max(promptTokens.write ?? 0, 0);
|
||||
const read = Math.max(promptTokens.read ?? 0, 0);
|
||||
const totalInputTokens = input + write + read;
|
||||
prompt = await createStructuredTransaction({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
inputTokens: -input,
|
||||
writeTokens: -write,
|
||||
readTokens: -read,
|
||||
inputTokenCount: totalInputTokens,
|
||||
});
|
||||
}
|
||||
|
||||
if (completionTokens) {
|
||||
const totalInputTokens = promptTokens
|
||||
? Math.max(promptTokens.input ?? 0, 0) +
|
||||
Math.max(promptTokens.write ?? 0, 0) +
|
||||
Math.max(promptTokens.read ?? 0, 0)
|
||||
: undefined;
|
||||
completion = await createTransaction({
|
||||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: -completionTokens,
|
||||
rawAmount: -Math.max(completionTokens, 0),
|
||||
inputTokenCount: totalInputTokens,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { createTransaction, createAutoRefillTransaction } = require('./Transaction');
|
||||
const { tokenValues, premiumTokenValues, getCacheMultiplier } = require('./tx');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
|
||||
require('~/db/models');
|
||||
|
||||
|
|
@ -734,4 +735,328 @@ describe('spendTokens', () => {
|
|||
expect(balance).toBeDefined();
|
||||
expect(balance.tokenCredits).toBeLessThan(10000); // Balance should be reduced
|
||||
});
|
||||
|
||||
describe('premium token pricing', () => {
|
||||
it('should charge standard rates for claude-opus-4-6 when prompt tokens are below threshold', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 100000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-standard-pricing',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const expectedCost =
|
||||
promptTokens * tokenValues[model].prompt + completionTokens * tokenValues[model].completion;
|
||||
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
it('should charge premium rates for claude-opus-4-6 when prompt tokens exceed threshold', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 250000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-premium-pricing',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const expectedCost =
|
||||
promptTokens * premiumTokenValues[model].prompt +
|
||||
completionTokens * premiumTokenValues[model].completion;
|
||||
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
it('should charge premium rates for both prompt and completion in structured tokens when above threshold', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-structured-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 200000,
|
||||
write: 10000,
|
||||
read: 5000,
|
||||
},
|
||||
completionTokens: 1000,
|
||||
};
|
||||
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
const writeRate = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readRate = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * premiumPromptRate +
|
||||
tokenUsage.promptTokens.write * writeRate +
|
||||
tokenUsage.promptTokens.read * readRate;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * premiumCompletionRate;
|
||||
|
||||
expect(result.prompt.prompt).toBeCloseTo(-expectedPromptCost, 0);
|
||||
expect(result.completion.completion).toBeCloseTo(-expectedCompletionCost, 0);
|
||||
});
|
||||
|
||||
it('should charge standard rates for structured tokens when below threshold', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-structured-standard',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 50000,
|
||||
write: 10000,
|
||||
read: 5000,
|
||||
},
|
||||
completionTokens: 1000,
|
||||
};
|
||||
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
const standardPromptRate = tokenValues[model].prompt;
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
const writeRate = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readRate = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * standardPromptRate +
|
||||
tokenUsage.promptTokens.write * writeRate +
|
||||
tokenUsage.promptTokens.read * readRate;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * standardCompletionRate;
|
||||
|
||||
expect(result.prompt.prompt).toBeCloseTo(-expectedPromptCost, 0);
|
||||
expect(result.completion.completion).toBeCloseTo(-expectedCompletionCost, 0);
|
||||
});
|
||||
|
||||
it('should not apply premium pricing to non-premium models regardless of prompt size', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-5';
|
||||
const promptTokens = 300000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-no-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const expectedCost =
|
||||
promptTokens * tokenValues[model].prompt + completionTokens * tokenValues[model].completion;
|
||||
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('inputTokenCount Normalization', () => {
|
||||
it('should normalize negative promptTokens to zero for inputTokenCount', async () => {
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: 100000000,
|
||||
});
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-negative-prompt',
|
||||
model: 'claude-opus-4-6',
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens: -500, completionTokens: 100 });
|
||||
|
||||
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
||||
|
||||
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
||||
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
||||
|
||||
expect(Math.abs(promptTx.rawAmount)).toBe(0);
|
||||
expect(completionTx.rawAmount).toBe(-100);
|
||||
|
||||
const standardCompletionRate = tokenValues['claude-opus-4-6'].completion;
|
||||
expect(completionTx.rate).toBe(standardCompletionRate);
|
||||
});
|
||||
|
||||
it('should use normalized inputTokenCount for premium threshold check on completion', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 250000;
|
||||
const completionTokens = 500;
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-normalized-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens, completionTokens });
|
||||
|
||||
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
||||
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
||||
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
||||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
expect(promptTx.rate).toBe(premiumPromptRate);
|
||||
expect(completionTx.rate).toBe(premiumCompletionRate);
|
||||
});
|
||||
|
||||
it('should keep inputTokenCount as zero when promptTokens is zero', async () => {
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: 100000000,
|
||||
});
|
||||
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-zero-prompt',
|
||||
model: 'claude-opus-4-6',
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens: 0, completionTokens: 100 });
|
||||
|
||||
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
||||
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
||||
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
||||
|
||||
expect(Math.abs(promptTx.rawAmount)).toBe(0);
|
||||
|
||||
const standardCompletionRate = tokenValues['claude-opus-4-6'].completion;
|
||||
expect(completionTx.rate).toBe(standardCompletionRate);
|
||||
});
|
||||
|
||||
it('should not trigger premium pricing with negative promptTokens on premium model', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-negative-no-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
await spendTokens(txData, { promptTokens: -300000, completionTokens: 500 });
|
||||
|
||||
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
||||
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
||||
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
expect(completionTx.rate).toBe(standardCompletionRate);
|
||||
});
|
||||
|
||||
it('should normalize negative structured token values to zero in spendStructuredTokens', async () => {
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({
|
||||
user: userId,
|
||||
tokenCredits: initialBalance,
|
||||
});
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-negative-structured',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: { input: -100, write: 50, read: -30 },
|
||||
completionTokens: -200,
|
||||
};
|
||||
|
||||
await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
const transactions = await Transaction.find({
|
||||
user: userId,
|
||||
conversationId: 'test-negative-structured',
|
||||
}).sort({ tokenType: 1 });
|
||||
|
||||
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
||||
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
||||
|
||||
expect(Math.abs(promptTx.inputTokens)).toBe(0);
|
||||
expect(promptTx.writeTokens).toBe(-50);
|
||||
expect(Math.abs(promptTx.readTokens)).toBe(0);
|
||||
|
||||
expect(Math.abs(completionTx.rawAmount)).toBe(0);
|
||||
|
||||
const standardRate = tokenValues[model].completion;
|
||||
expect(completionTx.rate).toBe(standardRate);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
132
api/models/tx.js
132
api/models/tx.js
|
|
@ -1,10 +1,40 @@
|
|||
const { matchModelName, findMatchingPattern } = require('@librechat/api');
|
||||
const defaultRate = 6;
|
||||
|
||||
/**
|
||||
* Token Pricing Configuration
|
||||
*
|
||||
* IMPORTANT: Key Ordering for Pattern Matching
|
||||
* ============================================
|
||||
* The `findMatchingPattern` function iterates through object keys in REVERSE order
|
||||
* (last-defined keys are checked first) and uses `modelName.includes(key)` for matching.
|
||||
*
|
||||
* This means:
|
||||
* 1. BASE PATTERNS must be defined FIRST (e.g., "kimi", "moonshot")
|
||||
* 2. SPECIFIC PATTERNS must be defined AFTER their base patterns (e.g., "kimi-k2", "kimi-k2.5")
|
||||
*
|
||||
* Example ordering for Kimi models:
|
||||
* kimi: { prompt: 0.6, completion: 2.5 }, // Base pattern - checked last
|
||||
* 'kimi-k2': { prompt: 0.6, completion: 2.5 }, // More specific - checked before "kimi"
|
||||
* 'kimi-k2.5': { prompt: 0.6, completion: 3.0 }, // Most specific - checked first
|
||||
*
|
||||
* Why this matters:
|
||||
* - Model name "kimi-k2.5" contains both "kimi" and "kimi-k2" as substrings
|
||||
* - If "kimi" were checked first, it would incorrectly match and return wrong pricing
|
||||
* - By defining specific patterns AFTER base patterns, they're checked first in reverse iteration
|
||||
*
|
||||
* This applies to BOTH `tokenValues` and `cacheTokenValues` objects.
|
||||
*
|
||||
* When adding new model families:
|
||||
* 1. Define the base/generic pattern first
|
||||
* 2. Define increasingly specific patterns after
|
||||
* 3. Ensure no pattern is a substring of another that should match differently
|
||||
*/
|
||||
|
||||
/**
|
||||
* AWS Bedrock pricing
|
||||
* source: https://aws.amazon.com/bedrock/pricing/
|
||||
* */
|
||||
*/
|
||||
const bedrockValues = {
|
||||
// Basic llama2 patterns (base defaults to smallest variant)
|
||||
llama2: { prompt: 0.75, completion: 1.0 },
|
||||
|
|
@ -80,6 +110,11 @@ const bedrockValues = {
|
|||
'nova-pro': { prompt: 0.8, completion: 3.2 },
|
||||
'nova-premier': { prompt: 2.5, completion: 12.5 },
|
||||
'deepseek.r1': { prompt: 1.35, completion: 5.4 },
|
||||
// Moonshot/Kimi models on Bedrock
|
||||
'moonshot.kimi': { prompt: 0.6, completion: 2.5 },
|
||||
'moonshot.kimi-k2': { prompt: 0.6, completion: 2.5 },
|
||||
'moonshot.kimi-k2.5': { prompt: 0.6, completion: 3.0 },
|
||||
'moonshot.kimi-k2-thinking': { prompt: 0.6, completion: 2.5 },
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -139,7 +174,9 @@ const tokenValues = Object.assign(
|
|||
'claude-haiku-4-5': { prompt: 1, completion: 5 },
|
||||
'claude-opus-4': { prompt: 15, completion: 75 },
|
||||
'claude-opus-4-5': { prompt: 5, completion: 25 },
|
||||
'claude-opus-4-6': { prompt: 5, completion: 25 },
|
||||
'claude-sonnet-4': { prompt: 3, completion: 15 },
|
||||
'claude-sonnet-4-6': { prompt: 3, completion: 15 },
|
||||
'command-r': { prompt: 0.5, completion: 1.5 },
|
||||
'command-r-plus': { prompt: 3, completion: 15 },
|
||||
'command-text': { prompt: 1.5, completion: 2.0 },
|
||||
|
|
@ -189,7 +226,31 @@ const tokenValues = Object.assign(
|
|||
'pixtral-large': { prompt: 2.0, completion: 6.0 },
|
||||
'mistral-large': { prompt: 2.0, completion: 6.0 },
|
||||
'mixtral-8x22b': { prompt: 0.65, completion: 0.65 },
|
||||
kimi: { prompt: 0.14, completion: 2.49 }, // Base pattern (using kimi-k2 pricing)
|
||||
// Moonshot/Kimi models (base patterns first, specific patterns last for correct matching)
|
||||
kimi: { prompt: 0.6, completion: 2.5 }, // Base pattern
|
||||
moonshot: { prompt: 2.0, completion: 5.0 }, // Base pattern (using 128k pricing)
|
||||
'kimi-latest': { prompt: 0.2, completion: 2.0 }, // Uses 8k/32k/128k pricing dynamically
|
||||
'kimi-k2': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2.5': { prompt: 0.6, completion: 3.0 },
|
||||
'kimi-k2-turbo': { prompt: 1.15, completion: 8.0 },
|
||||
'kimi-k2-turbo-preview': { prompt: 1.15, completion: 8.0 },
|
||||
'kimi-k2-0905': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2-0905-preview': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2-0711': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2-0711-preview': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2-thinking': { prompt: 0.6, completion: 2.5 },
|
||||
'kimi-k2-thinking-turbo': { prompt: 1.15, completion: 8.0 },
|
||||
'moonshot-v1': { prompt: 2.0, completion: 5.0 },
|
||||
'moonshot-v1-auto': { prompt: 2.0, completion: 5.0 },
|
||||
'moonshot-v1-8k': { prompt: 0.2, completion: 2.0 },
|
||||
'moonshot-v1-8k-vision': { prompt: 0.2, completion: 2.0 },
|
||||
'moonshot-v1-8k-vision-preview': { prompt: 0.2, completion: 2.0 },
|
||||
'moonshot-v1-32k': { prompt: 1.0, completion: 3.0 },
|
||||
'moonshot-v1-32k-vision': { prompt: 1.0, completion: 3.0 },
|
||||
'moonshot-v1-32k-vision-preview': { prompt: 1.0, completion: 3.0 },
|
||||
'moonshot-v1-128k': { prompt: 2.0, completion: 5.0 },
|
||||
'moonshot-v1-128k-vision': { prompt: 2.0, completion: 5.0 },
|
||||
'moonshot-v1-128k-vision-preview': { prompt: 2.0, completion: 5.0 },
|
||||
// GPT-OSS models (specific sizes)
|
||||
'gpt-oss:20b': { prompt: 0.05, completion: 0.2 },
|
||||
'gpt-oss-20b': { prompt: 0.05, completion: 0.2 },
|
||||
|
|
@ -249,12 +310,36 @@ const cacheTokenValues = {
|
|||
'claude-3-haiku': { write: 0.3, read: 0.03 },
|
||||
'claude-haiku-4-5': { write: 1.25, read: 0.1 },
|
||||
'claude-sonnet-4': { write: 3.75, read: 0.3 },
|
||||
'claude-sonnet-4-6': { write: 3.75, read: 0.3 },
|
||||
'claude-opus-4': { write: 18.75, read: 1.5 },
|
||||
'claude-opus-4-5': { write: 6.25, read: 0.5 },
|
||||
'claude-opus-4-6': { write: 6.25, read: 0.5 },
|
||||
// DeepSeek models - cache hit: $0.028/1M, cache miss: $0.28/1M
|
||||
deepseek: { write: 0.28, read: 0.028 },
|
||||
'deepseek-chat': { write: 0.28, read: 0.028 },
|
||||
'deepseek-reasoner': { write: 0.28, read: 0.028 },
|
||||
// Moonshot/Kimi models - cache hit: $0.15/1M (k2) or $0.10/1M (k2.5), cache miss: $0.60/1M
|
||||
kimi: { write: 0.6, read: 0.15 },
|
||||
'kimi-k2': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2.5': { write: 0.6, read: 0.1 },
|
||||
'kimi-k2-turbo': { write: 1.15, read: 0.15 },
|
||||
'kimi-k2-turbo-preview': { write: 1.15, read: 0.15 },
|
||||
'kimi-k2-0905': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2-0905-preview': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2-0711': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2-0711-preview': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2-thinking': { write: 0.6, read: 0.15 },
|
||||
'kimi-k2-thinking-turbo': { write: 1.15, read: 0.15 },
|
||||
};
|
||||
|
||||
/**
|
||||
* Premium (tiered) pricing for models whose rates change based on prompt size.
|
||||
* Each entry specifies the token threshold and the rates that apply above it.
|
||||
* @type {Object.<string, {threshold: number, prompt: number, completion: number}>}
|
||||
*/
|
||||
const premiumTokenValues = {
|
||||
'claude-opus-4-6': { threshold: 200000, prompt: 10, completion: 37.5 },
|
||||
'claude-sonnet-4-6': { threshold: 200000, prompt: 6, completion: 22.5 },
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -313,15 +398,27 @@ const getValueKey = (model, endpoint) => {
|
|||
* @param {string} [params.model] - The model name to derive the value key from if not provided.
|
||||
* @param {string} [params.endpoint] - The endpoint name to derive the value key from if not provided.
|
||||
* @param {EndpointTokenConfig} [params.endpointTokenConfig] - The token configuration for the endpoint.
|
||||
* @param {number} [params.inputTokenCount] - Total input token count for tiered pricing.
|
||||
* @returns {number} The multiplier for the given parameters, or a default value if not found.
|
||||
*/
|
||||
const getMultiplier = ({ valueKey, tokenType, model, endpoint, endpointTokenConfig }) => {
|
||||
const getMultiplier = ({
|
||||
model,
|
||||
valueKey,
|
||||
endpoint,
|
||||
tokenType,
|
||||
inputTokenCount,
|
||||
endpointTokenConfig,
|
||||
}) => {
|
||||
if (endpointTokenConfig) {
|
||||
return endpointTokenConfig?.[model]?.[tokenType] ?? defaultRate;
|
||||
}
|
||||
|
||||
if (valueKey && tokenType) {
|
||||
return tokenValues[valueKey][tokenType] ?? defaultRate;
|
||||
const premiumRate = getPremiumRate(valueKey, tokenType, inputTokenCount);
|
||||
if (premiumRate != null) {
|
||||
return premiumRate;
|
||||
}
|
||||
return tokenValues[valueKey]?.[tokenType] ?? defaultRate;
|
||||
}
|
||||
|
||||
if (!tokenType || !model) {
|
||||
|
|
@ -333,10 +430,33 @@ const getMultiplier = ({ valueKey, tokenType, model, endpoint, endpointTokenConf
|
|||
return defaultRate;
|
||||
}
|
||||
|
||||
// If we got this far, and values[tokenType] is undefined somehow, return a rough average of default multipliers
|
||||
const premiumRate = getPremiumRate(valueKey, tokenType, inputTokenCount);
|
||||
if (premiumRate != null) {
|
||||
return premiumRate;
|
||||
}
|
||||
|
||||
return tokenValues[valueKey]?.[tokenType] ?? defaultRate;
|
||||
};
|
||||
|
||||
/**
|
||||
* Checks if premium (tiered) pricing applies and returns the premium rate.
|
||||
* Each model defines its own threshold in `premiumTokenValues`.
|
||||
* @param {string} valueKey
|
||||
* @param {string} tokenType
|
||||
* @param {number} [inputTokenCount]
|
||||
* @returns {number|null}
|
||||
*/
|
||||
const getPremiumRate = (valueKey, tokenType, inputTokenCount) => {
|
||||
if (inputTokenCount == null) {
|
||||
return null;
|
||||
}
|
||||
const premiumEntry = premiumTokenValues[valueKey];
|
||||
if (!premiumEntry || inputTokenCount <= premiumEntry.threshold) {
|
||||
return null;
|
||||
}
|
||||
return premiumEntry[tokenType] ?? null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves the cache multiplier for a given value key and token type. If no value key is provided,
|
||||
* it attempts to derive it from the model name.
|
||||
|
|
@ -373,8 +493,10 @@ const getCacheMultiplier = ({ valueKey, cacheType, model, endpoint, endpointToke
|
|||
|
||||
module.exports = {
|
||||
tokenValues,
|
||||
premiumTokenValues,
|
||||
getValueKey,
|
||||
getMultiplier,
|
||||
getPremiumRate,
|
||||
getCacheMultiplier,
|
||||
defaultRate,
|
||||
cacheTokenValues,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
/** Note: No hard-coded values should be used in this file. */
|
||||
const { maxTokensMap } = require('@librechat/api');
|
||||
const { EModelEndpoint } = require('librechat-data-provider');
|
||||
const {
|
||||
|
|
@ -5,8 +6,10 @@ const {
|
|||
tokenValues,
|
||||
getValueKey,
|
||||
getMultiplier,
|
||||
getPremiumRate,
|
||||
cacheTokenValues,
|
||||
getCacheMultiplier,
|
||||
premiumTokenValues,
|
||||
} = require('./tx');
|
||||
|
||||
describe('getValueKey', () => {
|
||||
|
|
@ -239,6 +242,15 @@ describe('getMultiplier', () => {
|
|||
expect(getMultiplier({ valueKey: '8k', tokenType: 'unknownType' })).toBe(defaultRate);
|
||||
});
|
||||
|
||||
it('should return defaultRate if valueKey does not exist in tokenValues', () => {
|
||||
expect(getMultiplier({ valueKey: 'non-existent-model', tokenType: 'prompt' })).toBe(
|
||||
defaultRate,
|
||||
);
|
||||
expect(getMultiplier({ valueKey: 'non-existent-model', tokenType: 'completion' })).toBe(
|
||||
defaultRate,
|
||||
);
|
||||
});
|
||||
|
||||
it('should derive the valueKey from the model if not provided', () => {
|
||||
expect(getMultiplier({ tokenType: 'prompt', model: 'gpt-4-some-other-info' })).toBe(
|
||||
tokenValues['8k'].prompt,
|
||||
|
|
@ -334,8 +346,6 @@ describe('getMultiplier', () => {
|
|||
expect(getMultiplier({ model: 'openai/gpt-5.1', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['gpt-5.1'].prompt,
|
||||
);
|
||||
expect(tokenValues['gpt-5.1'].prompt).toBe(1.25);
|
||||
expect(tokenValues['gpt-5.1'].completion).toBe(10);
|
||||
});
|
||||
|
||||
it('should return the correct multiplier for gpt-5.2', () => {
|
||||
|
|
@ -348,8 +358,6 @@ describe('getMultiplier', () => {
|
|||
expect(getMultiplier({ model: 'openai/gpt-5.2', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['gpt-5.2'].prompt,
|
||||
);
|
||||
expect(tokenValues['gpt-5.2'].prompt).toBe(1.75);
|
||||
expect(tokenValues['gpt-5.2'].completion).toBe(14);
|
||||
});
|
||||
|
||||
it('should return the correct multiplier for gpt-4o', () => {
|
||||
|
|
@ -815,8 +823,6 @@ describe('Deepseek Model Tests', () => {
|
|||
expect(getMultiplier({ model: 'deepseek-chat', tokenType: 'completion' })).toBe(
|
||||
tokenValues['deepseek-chat'].completion,
|
||||
);
|
||||
expect(tokenValues['deepseek-chat'].prompt).toBe(0.28);
|
||||
expect(tokenValues['deepseek-chat'].completion).toBe(0.42);
|
||||
});
|
||||
|
||||
it('should return correct pricing for deepseek-reasoner', () => {
|
||||
|
|
@ -826,8 +832,6 @@ describe('Deepseek Model Tests', () => {
|
|||
expect(getMultiplier({ model: 'deepseek-reasoner', tokenType: 'completion' })).toBe(
|
||||
tokenValues['deepseek-reasoner'].completion,
|
||||
);
|
||||
expect(tokenValues['deepseek-reasoner'].prompt).toBe(0.28);
|
||||
expect(tokenValues['deepseek-reasoner'].completion).toBe(0.42);
|
||||
});
|
||||
|
||||
it('should handle DeepSeek model name variations with provider prefixes', () => {
|
||||
|
|
@ -840,8 +844,8 @@ describe('Deepseek Model Tests', () => {
|
|||
modelVariations.forEach((model) => {
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
expect(promptMultiplier).toBe(0.28);
|
||||
expect(completionMultiplier).toBe(0.42);
|
||||
expect(promptMultiplier).toBe(tokenValues['deepseek-chat'].prompt);
|
||||
expect(completionMultiplier).toBe(tokenValues['deepseek-chat'].completion);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -860,13 +864,13 @@ describe('Deepseek Model Tests', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('should return correct cache pricing values for DeepSeek models', () => {
|
||||
expect(cacheTokenValues['deepseek-chat'].write).toBe(0.28);
|
||||
expect(cacheTokenValues['deepseek-chat'].read).toBe(0.028);
|
||||
expect(cacheTokenValues['deepseek-reasoner'].write).toBe(0.28);
|
||||
expect(cacheTokenValues['deepseek-reasoner'].read).toBe(0.028);
|
||||
expect(cacheTokenValues['deepseek'].write).toBe(0.28);
|
||||
expect(cacheTokenValues['deepseek'].read).toBe(0.028);
|
||||
it('should have consistent cache pricing across DeepSeek model variants', () => {
|
||||
expect(cacheTokenValues['deepseek'].write).toBe(cacheTokenValues['deepseek-chat'].write);
|
||||
expect(cacheTokenValues['deepseek'].read).toBe(cacheTokenValues['deepseek-chat'].read);
|
||||
expect(cacheTokenValues['deepseek-reasoner'].write).toBe(
|
||||
cacheTokenValues['deepseek-chat'].write,
|
||||
);
|
||||
expect(cacheTokenValues['deepseek-reasoner'].read).toBe(cacheTokenValues['deepseek-chat'].read);
|
||||
});
|
||||
|
||||
it('should handle DeepSeek cache multipliers with model variations', () => {
|
||||
|
|
@ -875,8 +879,195 @@ describe('Deepseek Model Tests', () => {
|
|||
modelVariations.forEach((model) => {
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
expect(writeMultiplier).toBe(0.28);
|
||||
expect(readMultiplier).toBe(0.028);
|
||||
expect(writeMultiplier).toBe(cacheTokenValues['deepseek-chat'].write);
|
||||
expect(readMultiplier).toBe(cacheTokenValues['deepseek-chat'].read);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Moonshot/Kimi Model Tests - Pricing', () => {
|
||||
describe('Kimi Models', () => {
|
||||
it('should return correct pricing for kimi base pattern', () => {
|
||||
expect(getMultiplier({ model: 'kimi', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for kimi-k2.5', () => {
|
||||
expect(getMultiplier({ model: 'kimi-k2.5', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi-k2.5'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2.5', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi-k2.5'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for kimi-k2 series', () => {
|
||||
expect(getMultiplier({ model: 'kimi-k2', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi-k2'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi-k2'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for kimi-k2-turbo (higher pricing)', () => {
|
||||
expect(getMultiplier({ model: 'kimi-k2-turbo', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi-k2-turbo'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2-turbo', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi-k2-turbo'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for kimi-k2-thinking models', () => {
|
||||
expect(getMultiplier({ model: 'kimi-k2-thinking', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi-k2-thinking'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2-thinking', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi-k2-thinking'].completion,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2-thinking-turbo', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['kimi-k2-thinking-turbo'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'kimi-k2-thinking-turbo', tokenType: 'completion' })).toBe(
|
||||
tokenValues['kimi-k2-thinking-turbo'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle Kimi model variations with provider prefixes', () => {
|
||||
const modelVariations = ['openrouter/kimi-k2', 'openrouter/kimi-k2.5', 'openrouter/kimi'];
|
||||
|
||||
modelVariations.forEach((model) => {
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
expect(promptMultiplier).toBe(tokenValues['kimi'].prompt);
|
||||
expect([tokenValues['kimi'].completion, tokenValues['kimi-k2.5'].completion]).toContain(
|
||||
completionMultiplier,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Moonshot Models', () => {
|
||||
it('should return correct pricing for moonshot base pattern (128k pricing)', () => {
|
||||
expect(getMultiplier({ model: 'moonshot', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for moonshot-v1-8k', () => {
|
||||
expect(getMultiplier({ model: 'moonshot-v1-8k', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-8k'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-8k', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-8k'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for moonshot-v1-32k', () => {
|
||||
expect(getMultiplier({ model: 'moonshot-v1-32k', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-32k'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-32k', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-32k'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for moonshot-v1-128k', () => {
|
||||
expect(getMultiplier({ model: 'moonshot-v1-128k', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-128k'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-128k', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-128k'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct pricing for moonshot-v1 vision models', () => {
|
||||
expect(getMultiplier({ model: 'moonshot-v1-8k-vision', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-8k-vision'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-8k-vision', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-8k-vision'].completion,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-32k-vision', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-32k-vision'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-32k-vision', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-32k-vision'].completion,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-128k-vision', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot-v1-128k-vision'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot-v1-128k-vision', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot-v1-128k-vision'].completion,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Kimi Cache Multipliers', () => {
|
||||
it('should return correct cache multipliers for kimi-k2 models', () => {
|
||||
expect(getCacheMultiplier({ model: 'kimi', cacheType: 'write' })).toBe(
|
||||
cacheTokenValues['kimi'].write,
|
||||
);
|
||||
expect(getCacheMultiplier({ model: 'kimi', cacheType: 'read' })).toBe(
|
||||
cacheTokenValues['kimi'].read,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct cache multipliers for kimi-k2.5 (lower read price)', () => {
|
||||
expect(getCacheMultiplier({ model: 'kimi-k2.5', cacheType: 'write' })).toBe(
|
||||
cacheTokenValues['kimi-k2.5'].write,
|
||||
);
|
||||
expect(getCacheMultiplier({ model: 'kimi-k2.5', cacheType: 'read' })).toBe(
|
||||
cacheTokenValues['kimi-k2.5'].read,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return correct cache multipliers for kimi-k2-turbo', () => {
|
||||
expect(getCacheMultiplier({ model: 'kimi-k2-turbo', cacheType: 'write' })).toBe(
|
||||
cacheTokenValues['kimi-k2-turbo'].write,
|
||||
);
|
||||
expect(getCacheMultiplier({ model: 'kimi-k2-turbo', cacheType: 'read' })).toBe(
|
||||
cacheTokenValues['kimi-k2-turbo'].read,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle Kimi cache multipliers with model variations', () => {
|
||||
const modelVariations = ['openrouter/kimi-k2', 'openrouter/kimi'];
|
||||
|
||||
modelVariations.forEach((model) => {
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
expect(writeMultiplier).toBe(cacheTokenValues['kimi'].write);
|
||||
expect(readMultiplier).toBe(cacheTokenValues['kimi'].read);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Bedrock Moonshot Models', () => {
|
||||
it('should return correct pricing for Bedrock moonshot models', () => {
|
||||
expect(getMultiplier({ model: 'moonshot.kimi', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot.kimi'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot.kimi', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot.kimi'].completion,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot.kimi-k2', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot.kimi-k2'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot.kimi-k2.5', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['moonshot.kimi-k2.5'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'moonshot.kimi-k2.5', tokenType: 'completion' })).toBe(
|
||||
tokenValues['moonshot.kimi-k2.5'].completion,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1689,6 +1880,201 @@ describe('Claude Model Tests', () => {
|
|||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('should return correct prompt and completion rates for Claude Opus 4.6', () => {
|
||||
expect(getMultiplier({ model: 'claude-opus-4-6', tokenType: 'prompt' })).toBe(
|
||||
tokenValues['claude-opus-4-6'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: 'claude-opus-4-6', tokenType: 'completion' })).toBe(
|
||||
tokenValues['claude-opus-4-6'].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle Claude Opus 4.6 model name variations', () => {
|
||||
const modelVariations = [
|
||||
'claude-opus-4-6',
|
||||
'claude-opus-4-6-20250801',
|
||||
'claude-opus-4-6-latest',
|
||||
'anthropic/claude-opus-4-6',
|
||||
'claude-opus-4-6/anthropic',
|
||||
'claude-opus-4-6-preview',
|
||||
];
|
||||
|
||||
modelVariations.forEach((model) => {
|
||||
const valueKey = getValueKey(model);
|
||||
expect(valueKey).toBe('claude-opus-4-6');
|
||||
expect(getMultiplier({ model, tokenType: 'prompt' })).toBe(
|
||||
tokenValues['claude-opus-4-6'].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model, tokenType: 'completion' })).toBe(
|
||||
tokenValues['claude-opus-4-6'].completion,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('should return correct cache rates for Claude Opus 4.6', () => {
|
||||
expect(getCacheMultiplier({ model: 'claude-opus-4-6', cacheType: 'write' })).toBe(
|
||||
cacheTokenValues['claude-opus-4-6'].write,
|
||||
);
|
||||
expect(getCacheMultiplier({ model: 'claude-opus-4-6', cacheType: 'read' })).toBe(
|
||||
cacheTokenValues['claude-opus-4-6'].read,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle Claude Opus 4.6 cache rates with model name variations', () => {
|
||||
const modelVariations = [
|
||||
'claude-opus-4-6',
|
||||
'claude-opus-4-6-20250801',
|
||||
'claude-opus-4-6-latest',
|
||||
'anthropic/claude-opus-4-6',
|
||||
'claude-opus-4-6/anthropic',
|
||||
'claude-opus-4-6-preview',
|
||||
];
|
||||
|
||||
modelVariations.forEach((model) => {
|
||||
expect(getCacheMultiplier({ model, cacheType: 'write' })).toBe(
|
||||
cacheTokenValues['claude-opus-4-6'].write,
|
||||
);
|
||||
expect(getCacheMultiplier({ model, cacheType: 'read' })).toBe(
|
||||
cacheTokenValues['claude-opus-4-6'].read,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Premium Token Pricing', () => {
|
||||
const premiumModel = 'claude-opus-4-6';
|
||||
const premiumEntry = premiumTokenValues[premiumModel];
|
||||
const { threshold } = premiumEntry;
|
||||
const belowThreshold = threshold - 1;
|
||||
const aboveThreshold = threshold + 1;
|
||||
const wellAboveThreshold = threshold * 2;
|
||||
|
||||
it('should have premium pricing defined for claude-opus-4-6', () => {
|
||||
expect(premiumEntry).toBeDefined();
|
||||
expect(premiumEntry.threshold).toBeDefined();
|
||||
expect(premiumEntry.prompt).toBeDefined();
|
||||
expect(premiumEntry.completion).toBeDefined();
|
||||
expect(premiumEntry.prompt).toBeGreaterThan(tokenValues[premiumModel].prompt);
|
||||
expect(premiumEntry.completion).toBeGreaterThan(tokenValues[premiumModel].completion);
|
||||
});
|
||||
|
||||
it('should return null from getPremiumRate when inputTokenCount is below threshold', () => {
|
||||
expect(getPremiumRate(premiumModel, 'prompt', belowThreshold)).toBeNull();
|
||||
expect(getPremiumRate(premiumModel, 'completion', belowThreshold)).toBeNull();
|
||||
expect(getPremiumRate(premiumModel, 'prompt', threshold)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return premium rate from getPremiumRate when inputTokenCount exceeds threshold', () => {
|
||||
expect(getPremiumRate(premiumModel, 'prompt', aboveThreshold)).toBe(premiumEntry.prompt);
|
||||
expect(getPremiumRate(premiumModel, 'completion', aboveThreshold)).toBe(
|
||||
premiumEntry.completion,
|
||||
);
|
||||
expect(getPremiumRate(premiumModel, 'prompt', wellAboveThreshold)).toBe(premiumEntry.prompt);
|
||||
});
|
||||
|
||||
it('should return null from getPremiumRate when inputTokenCount is undefined or null', () => {
|
||||
expect(getPremiumRate(premiumModel, 'prompt', undefined)).toBeNull();
|
||||
expect(getPremiumRate(premiumModel, 'prompt', null)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null from getPremiumRate for models without premium pricing', () => {
|
||||
expect(getPremiumRate('claude-opus-4-5', 'prompt', wellAboveThreshold)).toBeNull();
|
||||
expect(getPremiumRate('claude-sonnet-4', 'prompt', wellAboveThreshold)).toBeNull();
|
||||
expect(getPremiumRate('gpt-4o', 'prompt', wellAboveThreshold)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return standard rate from getMultiplier when inputTokenCount is below threshold', () => {
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: premiumModel,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: belowThreshold,
|
||||
}),
|
||||
).toBe(tokenValues[premiumModel].prompt);
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: premiumModel,
|
||||
tokenType: 'completion',
|
||||
inputTokenCount: belowThreshold,
|
||||
}),
|
||||
).toBe(tokenValues[premiumModel].completion);
|
||||
});
|
||||
|
||||
it('should return premium rate from getMultiplier when inputTokenCount exceeds threshold', () => {
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: premiumModel,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: aboveThreshold,
|
||||
}),
|
||||
).toBe(premiumEntry.prompt);
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: premiumModel,
|
||||
tokenType: 'completion',
|
||||
inputTokenCount: aboveThreshold,
|
||||
}),
|
||||
).toBe(premiumEntry.completion);
|
||||
});
|
||||
|
||||
it('should return standard rate from getMultiplier when inputTokenCount is exactly at threshold', () => {
|
||||
expect(
|
||||
getMultiplier({ model: premiumModel, tokenType: 'prompt', inputTokenCount: threshold }),
|
||||
).toBe(tokenValues[premiumModel].prompt);
|
||||
});
|
||||
|
||||
it('should return premium rate from getMultiplier when inputTokenCount is one above threshold', () => {
|
||||
expect(
|
||||
getMultiplier({ model: premiumModel, tokenType: 'prompt', inputTokenCount: aboveThreshold }),
|
||||
).toBe(premiumEntry.prompt);
|
||||
});
|
||||
|
||||
it('should not apply premium pricing to models without premium entries', () => {
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: 'claude-opus-4-5',
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: wellAboveThreshold,
|
||||
}),
|
||||
).toBe(tokenValues['claude-opus-4-5'].prompt);
|
||||
expect(
|
||||
getMultiplier({
|
||||
model: 'claude-sonnet-4',
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: wellAboveThreshold,
|
||||
}),
|
||||
).toBe(tokenValues['claude-sonnet-4'].prompt);
|
||||
});
|
||||
|
||||
it('should use standard rate when inputTokenCount is not provided', () => {
|
||||
expect(getMultiplier({ model: premiumModel, tokenType: 'prompt' })).toBe(
|
||||
tokenValues[premiumModel].prompt,
|
||||
);
|
||||
expect(getMultiplier({ model: premiumModel, tokenType: 'completion' })).toBe(
|
||||
tokenValues[premiumModel].completion,
|
||||
);
|
||||
});
|
||||
|
||||
it('should apply premium pricing through getMultiplier with valueKey path', () => {
|
||||
const valueKey = getValueKey(premiumModel);
|
||||
expect(getMultiplier({ valueKey, tokenType: 'prompt', inputTokenCount: aboveThreshold })).toBe(
|
||||
premiumEntry.prompt,
|
||||
);
|
||||
expect(
|
||||
getMultiplier({ valueKey, tokenType: 'completion', inputTokenCount: aboveThreshold }),
|
||||
).toBe(premiumEntry.completion);
|
||||
});
|
||||
|
||||
it('should apply standard pricing through getMultiplier with valueKey path when below threshold', () => {
|
||||
const valueKey = getValueKey(premiumModel);
|
||||
expect(getMultiplier({ valueKey, tokenType: 'prompt', inputTokenCount: belowThreshold })).toBe(
|
||||
tokenValues[premiumModel].prompt,
|
||||
);
|
||||
expect(
|
||||
getMultiplier({ valueKey, tokenType: 'completion', inputTokenCount: belowThreshold }),
|
||||
).toBe(tokenValues[premiumModel].completion);
|
||||
});
|
||||
});
|
||||
|
||||
describe('tokens.ts and tx.js sync validation', () => {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"name": "@librechat/backend",
|
||||
"version": "v0.8.2-rc2",
|
||||
"version": "v0.8.2",
|
||||
"description": "",
|
||||
"scripts": {
|
||||
"start": "echo 'please run this from the root directory'",
|
||||
|
|
@ -34,25 +34,24 @@
|
|||
},
|
||||
"homepage": "https://librechat.ai",
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.71.0",
|
||||
"@anthropic-ai/vertex-sdk": "^0.14.0",
|
||||
"@aws-sdk/client-bedrock-runtime": "^3.941.0",
|
||||
"@aws-sdk/client-s3": "^3.758.0",
|
||||
"@anthropic-ai/vertex-sdk": "^0.14.3",
|
||||
"@aws-sdk/client-bedrock-runtime": "^3.980.0",
|
||||
"@aws-sdk/client-s3": "^3.980.0",
|
||||
"@aws-sdk/s3-request-presigner": "^3.758.0",
|
||||
"@azure/identity": "^4.7.0",
|
||||
"@azure/search-documents": "^12.0.0",
|
||||
"@azure/storage-blob": "^12.27.0",
|
||||
"@azure/storage-blob": "^12.30.0",
|
||||
"@google/genai": "^1.19.0",
|
||||
"@keyv/redis": "^4.3.3",
|
||||
"@langchain/core": "^0.3.80",
|
||||
"@librechat/agents": "^3.0.77",
|
||||
"@librechat/agents": "^3.1.50",
|
||||
"@librechat/api": "*",
|
||||
"@librechat/data-schemas": "*",
|
||||
"@microsoft/microsoft-graph-client": "^3.0.7",
|
||||
"@modelcontextprotocol/sdk": "^1.25.2",
|
||||
"@modelcontextprotocol/sdk": "^1.26.0",
|
||||
"@node-saml/passport-saml": "^5.1.0",
|
||||
"@smithy/node-http-handler": "^4.4.5",
|
||||
"axios": "^1.12.1",
|
||||
"axios": "^1.13.5",
|
||||
"bcryptjs": "^2.4.3",
|
||||
"compression": "^1.8.1",
|
||||
"connect-redis": "^8.1.0",
|
||||
|
|
@ -80,7 +79,7 @@
|
|||
"keyv-file": "^5.1.2",
|
||||
"klona": "^2.0.6",
|
||||
"librechat-data-provider": "*",
|
||||
"lodash": "^4.17.21",
|
||||
"lodash": "^4.17.23",
|
||||
"mathjs": "^15.1.0",
|
||||
"meilisearch": "^0.38.0",
|
||||
"memorystore": "^1.6.7",
|
||||
|
|
@ -108,7 +107,7 @@
|
|||
"tiktoken": "^1.0.15",
|
||||
"traverse": "^0.6.7",
|
||||
"ua-parser-js": "^1.0.36",
|
||||
"undici": "^7.10.0",
|
||||
"undici": "^7.18.2",
|
||||
"winston": "^3.11.0",
|
||||
"winston-daily-rotate-file": "^5.0.0",
|
||||
"zod": "^3.22.4"
|
||||
|
|
|
|||
|
|
@ -18,7 +18,6 @@ const {
|
|||
findUser,
|
||||
} = require('~/models');
|
||||
const { getGraphApiToken } = require('~/server/services/GraphTokenService');
|
||||
const { getOAuthReconnectionManager } = require('~/config');
|
||||
const { getOpenIdConfig } = require('~/strategies');
|
||||
|
||||
const registrationController = async (req, res) => {
|
||||
|
|
@ -79,7 +78,12 @@ const refreshController = async (req, res) => {
|
|||
|
||||
try {
|
||||
const openIdConfig = getOpenIdConfig();
|
||||
const tokenset = await openIdClient.refreshTokenGrant(openIdConfig, refreshToken);
|
||||
const refreshParams = process.env.OPENID_SCOPE ? { scope: process.env.OPENID_SCOPE } : {};
|
||||
const tokenset = await openIdClient.refreshTokenGrant(
|
||||
openIdConfig,
|
||||
refreshToken,
|
||||
refreshParams,
|
||||
);
|
||||
const claims = tokenset.claims();
|
||||
const { user, error, migration } = await findOpenIDUser({
|
||||
findUser,
|
||||
|
|
@ -161,17 +165,6 @@ const refreshController = async (req, res) => {
|
|||
if (session && session.expiration > new Date()) {
|
||||
const token = await setAuthTokens(userId, res, session);
|
||||
|
||||
// trigger OAuth MCP server reconnection asynchronously (best effort)
|
||||
try {
|
||||
void getOAuthReconnectionManager()
|
||||
.reconnectServers(userId)
|
||||
.catch((err) => {
|
||||
logger.error('[refreshController] Error reconnecting OAuth MCP servers:', err);
|
||||
});
|
||||
} catch (err) {
|
||||
logger.warn(`[refreshController] Cannot attempt OAuth MCP servers reconnection:`, err);
|
||||
}
|
||||
|
||||
res.status(200).send({ token, user });
|
||||
} else if (req?.query?.retry) {
|
||||
// Retrying from a refresh token request that failed (401)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { ResourceType, PrincipalType, PermissionBits } = require('librechat-data-provider');
|
||||
const { enrichRemoteAgentPrincipals, backfillRemoteAgentPermissions } = require('@librechat/api');
|
||||
const {
|
||||
bulkUpdateResourcePermissions,
|
||||
ensureGroupPrincipalExists,
|
||||
|
|
@ -14,7 +15,6 @@ const {
|
|||
findAccessibleResources,
|
||||
getResourcePermissionsMap,
|
||||
} = require('~/server/services/PermissionService');
|
||||
const { AclEntry } = require('~/db/models');
|
||||
const {
|
||||
searchPrincipals: searchLocalPrincipals,
|
||||
sortPrincipalsByRelevance,
|
||||
|
|
@ -24,6 +24,7 @@ const {
|
|||
entraIdPrincipalFeatureEnabled,
|
||||
searchEntraIdPrincipals,
|
||||
} = require('~/server/services/GraphApiService');
|
||||
const { AclEntry, AccessRole } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Generic controller for resource permission endpoints
|
||||
|
|
@ -234,7 +235,7 @@ const getResourcePermissions = async (req, res) => {
|
|||
},
|
||||
]);
|
||||
|
||||
const principals = [];
|
||||
let principals = [];
|
||||
let publicPermission = null;
|
||||
|
||||
// Process aggregation results
|
||||
|
|
@ -280,6 +281,13 @@ const getResourcePermissions = async (req, res) => {
|
|||
}
|
||||
}
|
||||
|
||||
if (resourceType === ResourceType.REMOTE_AGENT) {
|
||||
const enricherDeps = { AclEntry, AccessRole, logger };
|
||||
const enrichResult = await enrichRemoteAgentPrincipals(enricherDeps, resourceId, principals);
|
||||
principals = enrichResult.principals;
|
||||
backfillRemoteAgentPermissions(enricherDeps, resourceId, enrichResult.entriesToBackfill);
|
||||
}
|
||||
|
||||
// Return response in format expected by frontend
|
||||
const response = {
|
||||
resourceType,
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ const { getLogStores } = require('~/cache');
|
|||
|
||||
const getAvailablePluginsController = async (req, res) => {
|
||||
try {
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const cachedPlugins = await cache.get(CacheKeys.PLUGINS);
|
||||
if (cachedPlugins) {
|
||||
res.status(200).json(cachedPlugins);
|
||||
|
|
@ -63,7 +63,7 @@ const getAvailableTools = async (req, res) => {
|
|||
logger.warn('[getAvailableTools] User ID not found in request');
|
||||
return res.status(401).json({ message: 'Unauthorized' });
|
||||
}
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const cachedToolsArray = await cache.get(CacheKeys.TOOLS);
|
||||
|
||||
const appConfig = req.config ?? (await getAppConfig({ role: req.user?.role }));
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const { getCachedTools, getAppConfig } = require('~/server/services/Config');
|
||||
const { getLogStores } = require('~/cache');
|
||||
|
||||
|
|
@ -63,6 +64,28 @@ describe('PluginController', () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe('cache namespace', () => {
|
||||
it('getAvailablePluginsController should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue([]);
|
||||
await getAvailablePluginsController(mockReq, mockRes);
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
});
|
||||
|
||||
it('getAvailableTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue([]);
|
||||
await getAvailableTools(mockReq, mockRes);
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
});
|
||||
|
||||
it('should NOT use CONFIG_STORE namespace for tool/plugin operations', async () => {
|
||||
mockCache.get.mockResolvedValue([]);
|
||||
await getAvailablePluginsController(mockReq, mockRes);
|
||||
await getAvailableTools(mockReq, mockRes);
|
||||
const allCalls = getLogStores.mock.calls.flat();
|
||||
expect(allCalls).not.toContain(CacheKeys.CONFIG_STORE);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAvailablePluginsController', () => {
|
||||
it('should use filterUniquePlugins to remove duplicate plugins', async () => {
|
||||
// Add plugins with duplicates to availableTools
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ const {
|
|||
} = require('~/models');
|
||||
const {
|
||||
ConversationTag,
|
||||
AgentApiKey,
|
||||
Transaction,
|
||||
MemoryEntry,
|
||||
Assistant,
|
||||
|
|
@ -35,6 +36,7 @@ const {
|
|||
const { updateUserPluginAuth, deleteUserPluginAuth } = require('~/server/services/PluginService');
|
||||
const { verifyEmail, resendVerificationEmail } = require('~/server/services/AuthService');
|
||||
const { getMCPManager, getFlowStateManager, getMCPServersRegistry } = require('~/config');
|
||||
const { invalidateCachedTools } = require('~/server/services/Config/getCachedTools');
|
||||
const { needsRefresh, getNewS3URL } = require('~/server/services/Files/S3/crud');
|
||||
const { processDeleteRequest } = require('~/server/services/Files/process');
|
||||
const { getAppConfig } = require('~/server/services/Config');
|
||||
|
|
@ -214,6 +216,7 @@ const updateUserPluginsController = async (req, res) => {
|
|||
`[updateUserPluginsController] Attempting disconnect of MCP server "${serverName}" for user ${user.id} after plugin auth update.`,
|
||||
);
|
||||
await mcpManager.disconnectUserConnection(user.id, serverName);
|
||||
await invalidateCachedTools({ userId: user.id, serverName });
|
||||
}
|
||||
} catch (disconnectError) {
|
||||
logger.error(
|
||||
|
|
@ -256,6 +259,7 @@ const deleteUserController = async (req, res) => {
|
|||
await deleteFiles(null, user.id); // delete database files in case of orphaned files from previous steps
|
||||
await deleteToolCalls(user.id); // delete user tool calls
|
||||
await deleteUserAgents(user.id); // delete user agents
|
||||
await AgentApiKey.deleteMany({ user: user._id }); // delete user agent API keys
|
||||
await Assistant.deleteMany({ user: user.id }); // delete user assistants
|
||||
await ConversationTag.deleteMany({ user: user.id }); // delete user conversation tags
|
||||
await MemoryEntry.deleteMany({ userId: user.id }); // delete user memory entries
|
||||
|
|
|
|||
|
|
@ -16,13 +16,10 @@ jest.mock('@librechat/data-schemas', () => ({
|
|||
}));
|
||||
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
EnvVar: { CODE_API_KEY: 'CODE_API_KEY' },
|
||||
Providers: { GOOGLE: 'google' },
|
||||
GraphEvents: {},
|
||||
...jest.requireActual('@librechat/agents'),
|
||||
getMessageId: jest.fn(),
|
||||
ToolEndHandler: jest.fn(),
|
||||
handleToolCalls: jest.fn(),
|
||||
ChatModelStreamHandler: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/Citations', () => ({
|
||||
|
|
|
|||
281
api/server/controllers/agents/__tests__/jobReplacement.spec.js
Normal file
281
api/server/controllers/agents/__tests__/jobReplacement.spec.js
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
/**
|
||||
* Tests for job replacement detection in ResumableAgentController
|
||||
*
|
||||
* Tests the following fixes from PR #11462:
|
||||
* 1. Job creation timestamp tracking
|
||||
* 2. Stale job detection and event skipping
|
||||
* 3. Response message saving before final event emission
|
||||
*/
|
||||
|
||||
const mockLogger = {
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
};
|
||||
|
||||
const mockGenerationJobManager = {
|
||||
createJob: jest.fn(),
|
||||
getJob: jest.fn(),
|
||||
emitDone: jest.fn(),
|
||||
emitChunk: jest.fn(),
|
||||
completeJob: jest.fn(),
|
||||
updateMetadata: jest.fn(),
|
||||
setContentParts: jest.fn(),
|
||||
subscribe: jest.fn(),
|
||||
};
|
||||
|
||||
const mockSaveMessage = jest.fn();
|
||||
const mockDecrementPendingRequest = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: jest.fn().mockReturnValue(false),
|
||||
GenerationJobManager: mockGenerationJobManager,
|
||||
checkAndIncrementPendingRequest: jest.fn().mockResolvedValue({ allowed: true }),
|
||||
decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args),
|
||||
getViolationInfo: jest.fn(),
|
||||
sanitizeMessageForTransmit: jest.fn((msg) => msg),
|
||||
sanitizeFileForTransmit: jest.fn((file) => file),
|
||||
Constants: { NO_PARENT: '00000000-0000-0000-0000-000000000000' },
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: (...args) => mockSaveMessage(...args),
|
||||
}));
|
||||
|
||||
describe('Job Replacement Detection', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('Job Creation Timestamp Tracking', () => {
|
||||
it('should capture createdAt when job is created', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const createdAt = Date.now();
|
||||
|
||||
mockGenerationJobManager.createJob.mockResolvedValue({
|
||||
createdAt,
|
||||
readyPromise: Promise.resolve(),
|
||||
abortController: new AbortController(),
|
||||
emitter: { on: jest.fn() },
|
||||
});
|
||||
|
||||
const job = await mockGenerationJobManager.createJob(streamId, 'user-123', streamId);
|
||||
|
||||
expect(job.createdAt).toBe(createdAt);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Job Replacement Detection Logic', () => {
|
||||
/**
|
||||
* Simulates the job replacement detection logic from request.js
|
||||
* This is extracted for unit testing since the full controller is complex
|
||||
*/
|
||||
const detectJobReplacement = async (streamId, originalCreatedAt) => {
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
return !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
};
|
||||
|
||||
it('should detect when job was replaced (different createdAt)', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const newCreatedAt = 2000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: newCreatedAt,
|
||||
});
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect when job was deleted', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue(null);
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(true);
|
||||
});
|
||||
|
||||
it('should not detect replacement when same job (same createdAt)', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: originalCreatedAt,
|
||||
});
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Event Emission Behavior', () => {
|
||||
/**
|
||||
* Simulates the final event emission logic from request.js
|
||||
*/
|
||||
const emitFinalEventIfNotReplaced = async ({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent,
|
||||
userId,
|
||||
}) => {
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
|
||||
if (jobWasReplaced) {
|
||||
mockLogger.debug('Skipping FINAL emit - job was replaced', {
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
currentCreatedAt: currentJob?.createdAt,
|
||||
});
|
||||
await mockDecrementPendingRequest(userId);
|
||||
return false;
|
||||
}
|
||||
|
||||
mockGenerationJobManager.emitDone(streamId, finalEvent);
|
||||
mockGenerationJobManager.completeJob(streamId);
|
||||
await mockDecrementPendingRequest(userId);
|
||||
return true;
|
||||
};
|
||||
|
||||
it('should skip emitting when job was replaced', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const newCreatedAt = 2000;
|
||||
const userId = 'user-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: newCreatedAt,
|
||||
});
|
||||
|
||||
const emitted = await emitFinalEventIfNotReplaced({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent: { final: true },
|
||||
userId,
|
||||
});
|
||||
|
||||
expect(emitted).toBe(false);
|
||||
expect(mockGenerationJobManager.emitDone).not.toHaveBeenCalled();
|
||||
expect(mockGenerationJobManager.completeJob).not.toHaveBeenCalled();
|
||||
expect(mockDecrementPendingRequest).toHaveBeenCalledWith(userId);
|
||||
expect(mockLogger.debug).toHaveBeenCalledWith(
|
||||
'Skipping FINAL emit - job was replaced',
|
||||
expect.objectContaining({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
currentCreatedAt: newCreatedAt,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should emit when job was not replaced', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const userId = 'user-123';
|
||||
const finalEvent = { final: true, conversation: { conversationId: streamId } };
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: originalCreatedAt,
|
||||
});
|
||||
|
||||
const emitted = await emitFinalEventIfNotReplaced({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent,
|
||||
userId,
|
||||
});
|
||||
|
||||
expect(emitted).toBe(true);
|
||||
expect(mockGenerationJobManager.emitDone).toHaveBeenCalledWith(streamId, finalEvent);
|
||||
expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(streamId);
|
||||
expect(mockDecrementPendingRequest).toHaveBeenCalledWith(userId);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Response Message Saving Order', () => {
|
||||
/**
|
||||
* Tests that response messages are saved BEFORE final events are emitted
|
||||
* This prevents race conditions where clients send follow-up messages
|
||||
* before the response is in the database
|
||||
*/
|
||||
it('should save message before emitting final event', async () => {
|
||||
const callOrder = [];
|
||||
|
||||
mockSaveMessage.mockImplementation(async () => {
|
||||
callOrder.push('saveMessage');
|
||||
});
|
||||
|
||||
mockGenerationJobManager.emitDone.mockImplementation(() => {
|
||||
callOrder.push('emitDone');
|
||||
});
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: 1000,
|
||||
});
|
||||
|
||||
// Simulate the order of operations from request.js
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const response = { messageId: 'response-123' };
|
||||
const userId = 'user-123';
|
||||
|
||||
// Step 1: Save message
|
||||
await mockSaveMessage({}, { ...response, user: userId }, { context: 'test' });
|
||||
|
||||
// Step 2: Check for replacement
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
|
||||
// Step 3: Emit if not replaced
|
||||
if (!jobWasReplaced) {
|
||||
mockGenerationJobManager.emitDone(streamId, { final: true });
|
||||
}
|
||||
|
||||
expect(callOrder).toEqual(['saveMessage', 'emitDone']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Aborted Request Handling', () => {
|
||||
it('should use unfinished: true instead of error: true for aborted requests', () => {
|
||||
const response = { messageId: 'response-123', content: [] };
|
||||
|
||||
// The new format for aborted responses
|
||||
const abortedResponse = { ...response, unfinished: true };
|
||||
|
||||
expect(abortedResponse.unfinished).toBe(true);
|
||||
expect(abortedResponse.error).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should include unfinished flag in final event for aborted requests', () => {
|
||||
const response = { messageId: 'response-123', content: [] };
|
||||
|
||||
// Old format (deprecated)
|
||||
const _oldFinalEvent = {
|
||||
final: true,
|
||||
responseMessage: { ...response, error: true },
|
||||
error: { message: 'Request was aborted' },
|
||||
};
|
||||
|
||||
// New format (PR #11462)
|
||||
const newFinalEvent = {
|
||||
final: true,
|
||||
responseMessage: { ...response, unfinished: true },
|
||||
};
|
||||
|
||||
expect(newFinalEvent.responseMessage.unfinished).toBe(true);
|
||||
expect(newFinalEvent.error).toBeUndefined();
|
||||
expect(newFinalEvent.responseMessage.error).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
204
api/server/controllers/agents/__tests__/openai.spec.js
Normal file
204
api/server/controllers/agents/__tests__/openai.spec.js
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
/**
|
||||
* Unit tests for OpenAI-compatible API controller
|
||||
* Tests that recordCollectedUsage is called correctly for token spending
|
||||
*/
|
||||
|
||||
const mockSpendTokens = jest.fn().mockResolvedValue({});
|
||||
const mockSpendStructuredTokens = jest.fn().mockResolvedValue({});
|
||||
const mockRecordCollectedUsage = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ input_tokens: 100, output_tokens: 50 });
|
||||
const mockGetBalanceConfig = jest.fn().mockReturnValue({ enabled: true });
|
||||
const mockGetTransactionsConfig = jest.fn().mockReturnValue({ enabled: true });
|
||||
|
||||
jest.mock('nanoid', () => ({
|
||||
nanoid: jest.fn(() => 'mock-nanoid-123'),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
Callback: { TOOL_ERROR: 'TOOL_ERROR' },
|
||||
ToolEndHandler: jest.fn(),
|
||||
formatAgentMessages: jest.fn().mockReturnValue({
|
||||
messages: [],
|
||||
indexTokenCountMap: {},
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
writeSSE: jest.fn(),
|
||||
createRun: jest.fn().mockResolvedValue({
|
||||
processStream: jest.fn().mockResolvedValue(undefined),
|
||||
}),
|
||||
createChunk: jest.fn().mockReturnValue({}),
|
||||
buildToolSet: jest.fn().mockReturnValue(new Set()),
|
||||
sendFinalChunk: jest.fn(),
|
||||
createSafeUser: jest.fn().mockReturnValue({ id: 'user-123' }),
|
||||
validateRequest: jest
|
||||
.fn()
|
||||
.mockReturnValue({ request: { model: 'agent-123', messages: [], stream: false } }),
|
||||
initializeAgent: jest.fn().mockResolvedValue({
|
||||
model: 'gpt-4',
|
||||
model_parameters: {},
|
||||
toolRegistry: {},
|
||||
}),
|
||||
getBalanceConfig: mockGetBalanceConfig,
|
||||
createErrorResponse: jest.fn(),
|
||||
getTransactionsConfig: mockGetTransactionsConfig,
|
||||
recordCollectedUsage: mockRecordCollectedUsage,
|
||||
buildNonStreamingResponse: jest.fn().mockReturnValue({ id: 'resp-123' }),
|
||||
createOpenAIStreamTracker: jest.fn().mockReturnValue({
|
||||
addText: jest.fn(),
|
||||
addReasoning: jest.fn(),
|
||||
toolCalls: new Map(),
|
||||
usage: { promptTokens: 0, completionTokens: 0, reasoningTokens: 0 },
|
||||
}),
|
||||
createOpenAIContentAggregator: jest.fn().mockReturnValue({
|
||||
addText: jest.fn(),
|
||||
addReasoning: jest.fn(),
|
||||
getText: jest.fn().mockReturnValue(''),
|
||||
getReasoning: jest.fn().mockReturnValue(''),
|
||||
toolCalls: new Map(),
|
||||
usage: { promptTokens: 100, completionTokens: 50, reasoningTokens: 0 },
|
||||
}),
|
||||
createToolExecuteHandler: jest.fn().mockReturnValue({ handle: jest.fn() }),
|
||||
isChatCompletionValidationFailure: jest.fn().mockReturnValue(false),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/ToolService', () => ({
|
||||
loadAgentTools: jest.fn().mockResolvedValue([]),
|
||||
loadToolsForExecution: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/spendTokens', () => ({
|
||||
spendTokens: mockSpendTokens,
|
||||
spendStructuredTokens: mockSpendStructuredTokens,
|
||||
}));
|
||||
|
||||
jest.mock('~/server/controllers/agents/callbacks', () => ({
|
||||
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/PermissionService', () => ({
|
||||
findAccessibleResources: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Conversation', () => ({
|
||||
getConvoFiles: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Agent', () => ({
|
||||
getAgent: jest.fn().mockResolvedValue({
|
||||
id: 'agent-123',
|
||||
provider: 'openAI',
|
||||
model_parameters: { model: 'gpt-4' },
|
||||
}),
|
||||
getAgents: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
getFiles: jest.fn(),
|
||||
getUserKey: jest.fn(),
|
||||
getMessages: jest.fn(),
|
||||
updateFilesUsage: jest.fn(),
|
||||
getUserKeyValues: jest.fn(),
|
||||
getUserCodeFiles: jest.fn(),
|
||||
getToolFilesByIds: jest.fn(),
|
||||
getCodeGeneratedFiles: jest.fn(),
|
||||
}));
|
||||
|
||||
describe('OpenAIChatCompletionController', () => {
|
||||
let OpenAIChatCompletionController;
|
||||
let req, res;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
|
||||
const controller = require('../openai');
|
||||
OpenAIChatCompletionController = controller.OpenAIChatCompletionController;
|
||||
|
||||
req = {
|
||||
body: {
|
||||
model: 'agent-123',
|
||||
messages: [{ role: 'user', content: 'Hello' }],
|
||||
stream: false,
|
||||
},
|
||||
user: { id: 'user-123' },
|
||||
config: {
|
||||
endpoints: {
|
||||
agents: { allowedProviders: ['openAI'] },
|
||||
},
|
||||
},
|
||||
on: jest.fn(),
|
||||
};
|
||||
|
||||
res = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
setHeader: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
end: jest.fn(),
|
||||
write: jest.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
describe('token usage recording', () => {
|
||||
it('should call recordCollectedUsage after successful non-streaming completion', async () => {
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
{ spendTokens: mockSpendTokens, spendStructuredTokens: mockSpendStructuredTokens },
|
||||
expect.objectContaining({
|
||||
user: 'user-123',
|
||||
conversationId: expect.any(String),
|
||||
collectedUsage: expect.any(Array),
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass balance and transactions config to recordCollectedUsage', async () => {
|
||||
mockGetBalanceConfig.mockReturnValue({ enabled: true, startBalance: 1000 });
|
||||
mockGetTransactionsConfig.mockReturnValue({ enabled: true, rateLimit: 100 });
|
||||
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
balance: { enabled: true, startBalance: 1000 },
|
||||
transactions: { enabled: true, rateLimit: 100 },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass spendTokens and spendStructuredTokens as dependencies', async () => {
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
const [deps] = mockRecordCollectedUsage.mock.calls[0];
|
||||
expect(deps).toHaveProperty('spendTokens', mockSpendTokens);
|
||||
expect(deps).toHaveProperty('spendStructuredTokens', mockSpendStructuredTokens);
|
||||
});
|
||||
|
||||
it('should include model from primaryConfig in recordCollectedUsage params', async () => {
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
model: 'gpt-4',
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
312
api/server/controllers/agents/__tests__/responses.unit.spec.js
Normal file
312
api/server/controllers/agents/__tests__/responses.unit.spec.js
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
/**
|
||||
* Unit tests for Open Responses API controller
|
||||
* Tests that recordCollectedUsage is called correctly for token spending
|
||||
*/
|
||||
|
||||
const mockSpendTokens = jest.fn().mockResolvedValue({});
|
||||
const mockSpendStructuredTokens = jest.fn().mockResolvedValue({});
|
||||
const mockRecordCollectedUsage = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ input_tokens: 100, output_tokens: 50 });
|
||||
const mockGetBalanceConfig = jest.fn().mockReturnValue({ enabled: true });
|
||||
const mockGetTransactionsConfig = jest.fn().mockReturnValue({ enabled: true });
|
||||
|
||||
jest.mock('nanoid', () => ({
|
||||
nanoid: jest.fn(() => 'mock-nanoid-123'),
|
||||
}));
|
||||
|
||||
jest.mock('uuid', () => ({
|
||||
v4: jest.fn(() => 'mock-uuid-456'),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
Callback: { TOOL_ERROR: 'TOOL_ERROR' },
|
||||
ToolEndHandler: jest.fn(),
|
||||
formatAgentMessages: jest.fn().mockReturnValue({
|
||||
messages: [],
|
||||
indexTokenCountMap: {},
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
createRun: jest.fn().mockResolvedValue({
|
||||
processStream: jest.fn().mockResolvedValue(undefined),
|
||||
}),
|
||||
buildToolSet: jest.fn().mockReturnValue(new Set()),
|
||||
createSafeUser: jest.fn().mockReturnValue({ id: 'user-123' }),
|
||||
initializeAgent: jest.fn().mockResolvedValue({
|
||||
model: 'claude-3',
|
||||
model_parameters: {},
|
||||
toolRegistry: {},
|
||||
}),
|
||||
getBalanceConfig: mockGetBalanceConfig,
|
||||
getTransactionsConfig: mockGetTransactionsConfig,
|
||||
recordCollectedUsage: mockRecordCollectedUsage,
|
||||
createToolExecuteHandler: jest.fn().mockReturnValue({ handle: jest.fn() }),
|
||||
// Responses API
|
||||
writeDone: jest.fn(),
|
||||
buildResponse: jest.fn().mockReturnValue({ id: 'resp_123', output: [] }),
|
||||
generateResponseId: jest.fn().mockReturnValue('resp_mock-123'),
|
||||
isValidationFailure: jest.fn().mockReturnValue(false),
|
||||
emitResponseCreated: jest.fn(),
|
||||
createResponseContext: jest.fn().mockReturnValue({ responseId: 'resp_123' }),
|
||||
createResponseTracker: jest.fn().mockReturnValue({
|
||||
usage: { promptTokens: 100, completionTokens: 50 },
|
||||
}),
|
||||
setupStreamingResponse: jest.fn(),
|
||||
emitResponseInProgress: jest.fn(),
|
||||
convertInputToMessages: jest.fn().mockReturnValue([]),
|
||||
validateResponseRequest: jest.fn().mockReturnValue({
|
||||
request: { model: 'agent-123', input: 'Hello', stream: false },
|
||||
}),
|
||||
buildAggregatedResponse: jest.fn().mockReturnValue({
|
||||
id: 'resp_123',
|
||||
status: 'completed',
|
||||
output: [],
|
||||
usage: { input_tokens: 100, output_tokens: 50, total_tokens: 150 },
|
||||
}),
|
||||
createResponseAggregator: jest.fn().mockReturnValue({
|
||||
usage: { promptTokens: 100, completionTokens: 50 },
|
||||
}),
|
||||
sendResponsesErrorResponse: jest.fn(),
|
||||
createResponsesEventHandlers: jest.fn().mockReturnValue({
|
||||
handlers: {
|
||||
on_message_delta: { handle: jest.fn() },
|
||||
on_reasoning_delta: { handle: jest.fn() },
|
||||
on_run_step: { handle: jest.fn() },
|
||||
on_run_step_delta: { handle: jest.fn() },
|
||||
on_chat_model_end: { handle: jest.fn() },
|
||||
},
|
||||
finalizeStream: jest.fn(),
|
||||
}),
|
||||
createAggregatorEventHandlers: jest.fn().mockReturnValue({
|
||||
on_message_delta: { handle: jest.fn() },
|
||||
on_reasoning_delta: { handle: jest.fn() },
|
||||
on_run_step: { handle: jest.fn() },
|
||||
on_run_step_delta: { handle: jest.fn() },
|
||||
on_chat_model_end: { handle: jest.fn() },
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/ToolService', () => ({
|
||||
loadAgentTools: jest.fn().mockResolvedValue([]),
|
||||
loadToolsForExecution: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/spendTokens', () => ({
|
||||
spendTokens: mockSpendTokens,
|
||||
spendStructuredTokens: mockSpendStructuredTokens,
|
||||
}));
|
||||
|
||||
jest.mock('~/server/controllers/agents/callbacks', () => ({
|
||||
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
createResponsesToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/PermissionService', () => ({
|
||||
findAccessibleResources: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Conversation', () => ({
|
||||
getConvoFiles: jest.fn().mockResolvedValue([]),
|
||||
saveConvo: jest.fn().mockResolvedValue({}),
|
||||
getConvo: jest.fn().mockResolvedValue(null),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Agent', () => ({
|
||||
getAgent: jest.fn().mockResolvedValue({
|
||||
id: 'agent-123',
|
||||
name: 'Test Agent',
|
||||
provider: 'anthropic',
|
||||
model_parameters: { model: 'claude-3' },
|
||||
}),
|
||||
getAgents: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
getFiles: jest.fn(),
|
||||
getUserKey: jest.fn(),
|
||||
getMessages: jest.fn().mockResolvedValue([]),
|
||||
saveMessage: jest.fn().mockResolvedValue({}),
|
||||
updateFilesUsage: jest.fn(),
|
||||
getUserKeyValues: jest.fn(),
|
||||
getUserCodeFiles: jest.fn(),
|
||||
getToolFilesByIds: jest.fn(),
|
||||
getCodeGeneratedFiles: jest.fn(),
|
||||
}));
|
||||
|
||||
describe('createResponse controller', () => {
|
||||
let createResponse;
|
||||
let req, res;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
|
||||
const controller = require('../responses');
|
||||
createResponse = controller.createResponse;
|
||||
|
||||
req = {
|
||||
body: {
|
||||
model: 'agent-123',
|
||||
input: 'Hello',
|
||||
stream: false,
|
||||
},
|
||||
user: { id: 'user-123' },
|
||||
config: {
|
||||
endpoints: {
|
||||
agents: { allowedProviders: ['anthropic'] },
|
||||
},
|
||||
},
|
||||
on: jest.fn(),
|
||||
};
|
||||
|
||||
res = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
setHeader: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
end: jest.fn(),
|
||||
write: jest.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
describe('token usage recording - non-streaming', () => {
|
||||
it('should call recordCollectedUsage after successful non-streaming completion', async () => {
|
||||
await createResponse(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
{ spendTokens: mockSpendTokens, spendStructuredTokens: mockSpendStructuredTokens },
|
||||
expect.objectContaining({
|
||||
user: 'user-123',
|
||||
conversationId: expect.any(String),
|
||||
collectedUsage: expect.any(Array),
|
||||
context: 'message',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass balance and transactions config to recordCollectedUsage', async () => {
|
||||
mockGetBalanceConfig.mockReturnValue({ enabled: true, startBalance: 2000 });
|
||||
mockGetTransactionsConfig.mockReturnValue({ enabled: true });
|
||||
|
||||
await createResponse(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
balance: { enabled: true, startBalance: 2000 },
|
||||
transactions: { enabled: true },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass spendTokens and spendStructuredTokens as dependencies', async () => {
|
||||
await createResponse(req, res);
|
||||
|
||||
const [deps] = mockRecordCollectedUsage.mock.calls[0];
|
||||
expect(deps).toHaveProperty('spendTokens', mockSpendTokens);
|
||||
expect(deps).toHaveProperty('spendStructuredTokens', mockSpendStructuredTokens);
|
||||
});
|
||||
|
||||
it('should include model from primaryConfig in recordCollectedUsage params', async () => {
|
||||
await createResponse(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
model: 'claude-3',
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('token usage recording - streaming', () => {
|
||||
beforeEach(() => {
|
||||
req.body.stream = true;
|
||||
|
||||
const api = require('@librechat/api');
|
||||
api.validateResponseRequest.mockReturnValue({
|
||||
request: { model: 'agent-123', input: 'Hello', stream: true },
|
||||
});
|
||||
});
|
||||
|
||||
it('should call recordCollectedUsage after successful streaming completion', async () => {
|
||||
await createResponse(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
{ spendTokens: mockSpendTokens, spendStructuredTokens: mockSpendStructuredTokens },
|
||||
expect.objectContaining({
|
||||
user: 'user-123',
|
||||
context: 'message',
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('collectedUsage population', () => {
|
||||
it('should collect usage from on_chat_model_end events', async () => {
|
||||
const api = require('@librechat/api');
|
||||
|
||||
let capturedOnChatModelEnd;
|
||||
api.createAggregatorEventHandlers.mockImplementation(() => {
|
||||
return {
|
||||
on_message_delta: { handle: jest.fn() },
|
||||
on_reasoning_delta: { handle: jest.fn() },
|
||||
on_run_step: { handle: jest.fn() },
|
||||
on_run_step_delta: { handle: jest.fn() },
|
||||
on_chat_model_end: {
|
||||
handle: jest.fn((event, data) => {
|
||||
if (capturedOnChatModelEnd) {
|
||||
capturedOnChatModelEnd(event, data);
|
||||
}
|
||||
}),
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
api.createRun.mockImplementation(async ({ customHandlers }) => {
|
||||
capturedOnChatModelEnd = (event, data) => {
|
||||
customHandlers.on_chat_model_end.handle(event, data);
|
||||
};
|
||||
|
||||
return {
|
||||
processStream: jest.fn().mockImplementation(async () => {
|
||||
customHandlers.on_chat_model_end.handle('on_chat_model_end', {
|
||||
output: {
|
||||
usage_metadata: {
|
||||
input_tokens: 150,
|
||||
output_tokens: 75,
|
||||
model: 'claude-3',
|
||||
},
|
||||
},
|
||||
});
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
await createResponse(req, res);
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
collectedUsage: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
input_tokens: 150,
|
||||
output_tokens: 75,
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,16 +1,13 @@
|
|||
const { nanoid } = require('nanoid');
|
||||
const { sendEvent, GenerationJobManager } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Constants, EnvVar, GraphEvents, ToolEndHandler } = require('@librechat/agents');
|
||||
const { Tools, StepTypes, FileContext, ErrorTypes } = require('librechat-data-provider');
|
||||
const {
|
||||
EnvVar,
|
||||
Providers,
|
||||
GraphEvents,
|
||||
getMessageId,
|
||||
ToolEndHandler,
|
||||
handleToolCalls,
|
||||
ChatModelStreamHandler,
|
||||
} = require('@librechat/agents');
|
||||
sendEvent,
|
||||
GenerationJobManager,
|
||||
writeAttachmentEvent,
|
||||
createToolExecuteHandler,
|
||||
} = require('@librechat/api');
|
||||
const { processFileCitations } = require('~/server/services/Files/Citations');
|
||||
const { processCodeOutput } = require('~/server/services/Files/Code/process');
|
||||
const { loadAuthValues } = require('~/server/services/Tools/credentials');
|
||||
|
|
@ -51,8 +48,6 @@ class ModelEndHandler {
|
|||
let errorMessage;
|
||||
try {
|
||||
const agentContext = graph.getAgentContext(metadata);
|
||||
const isGoogle = agentContext.provider === Providers.GOOGLE;
|
||||
const streamingDisabled = !!agentContext.clientOptions?.disableStreaming;
|
||||
if (data?.output?.additional_kwargs?.stop_reason === 'refusal') {
|
||||
const info = { ...data.output.additional_kwargs };
|
||||
errorMessage = JSON.stringify({
|
||||
|
|
@ -67,21 +62,6 @@ class ModelEndHandler {
|
|||
});
|
||||
}
|
||||
|
||||
const toolCalls = data?.output?.tool_calls;
|
||||
let hasUnprocessedToolCalls = false;
|
||||
if (Array.isArray(toolCalls) && toolCalls.length > 0 && graph?.toolCallStepIds?.has) {
|
||||
try {
|
||||
hasUnprocessedToolCalls = toolCalls.some(
|
||||
(tc) => tc?.id && !graph.toolCallStepIds.has(tc.id),
|
||||
);
|
||||
} catch {
|
||||
hasUnprocessedToolCalls = false;
|
||||
}
|
||||
}
|
||||
if (isGoogle || streamingDisabled || hasUnprocessedToolCalls) {
|
||||
await handleToolCalls(toolCalls, metadata, graph);
|
||||
}
|
||||
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (!usage) {
|
||||
return this.finalize(errorMessage);
|
||||
|
|
@ -92,38 +72,6 @@ class ModelEndHandler {
|
|||
}
|
||||
|
||||
this.collectedUsage.push(usage);
|
||||
if (!streamingDisabled) {
|
||||
return this.finalize(errorMessage);
|
||||
}
|
||||
if (!data.output.content) {
|
||||
return this.finalize(errorMessage);
|
||||
}
|
||||
const stepKey = graph.getStepKey(metadata);
|
||||
const message_id = getMessageId(stepKey, graph) ?? '';
|
||||
if (message_id) {
|
||||
await graph.dispatchRunStep(stepKey, {
|
||||
type: StepTypes.MESSAGE_CREATION,
|
||||
message_creation: {
|
||||
message_id,
|
||||
},
|
||||
});
|
||||
}
|
||||
const stepId = graph.getStepIdByKey(stepKey);
|
||||
const content = data.output.content;
|
||||
if (typeof content === 'string') {
|
||||
await graph.dispatchMessageDelta(stepId, {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: content,
|
||||
},
|
||||
],
|
||||
});
|
||||
} else if (content.every((c) => c.type?.startsWith('text'))) {
|
||||
await graph.dispatchMessageDelta(stepId, {
|
||||
content,
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error handling model end event:', error);
|
||||
return this.finalize(errorMessage);
|
||||
|
|
@ -146,18 +94,26 @@ function checkIfLastAgent(last_agent_id, langgraph_node) {
|
|||
|
||||
/**
|
||||
* Helper to emit events either to res (standard mode) or to job emitter (resumable mode).
|
||||
* In Redis mode, awaits the emit to guarantee event ordering (critical for streaming deltas).
|
||||
* @param {ServerResponse} res - The server response object
|
||||
* @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode
|
||||
* @param {Object} eventData - The event data to send
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
function emitEvent(res, streamId, eventData) {
|
||||
async function emitEvent(res, streamId, eventData) {
|
||||
if (streamId) {
|
||||
GenerationJobManager.emitChunk(streamId, eventData);
|
||||
await GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
sendEvent(res, eventData);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @typedef {Object} ToolExecuteOptions
|
||||
* @property {(toolNames: string[]) => Promise<{loadedTools: StructuredTool[]}>} loadTools - Function to load tools by name
|
||||
* @property {Object} configurable - Configurable context for tool invocation
|
||||
*/
|
||||
|
||||
/**
|
||||
* Get default handlers for stream events.
|
||||
* @param {Object} options - The options object.
|
||||
|
|
@ -166,6 +122,7 @@ function emitEvent(res, streamId, eventData) {
|
|||
* @param {ToolEndCallback} options.toolEndCallback - Callback to use when tool ends.
|
||||
* @param {Array<UsageMetadata>} options.collectedUsage - The list of collected usage metadata.
|
||||
* @param {string | null} [options.streamId] - The stream ID for resumable mode, or null for standard mode.
|
||||
* @param {ToolExecuteOptions} [options.toolExecuteOptions] - Options for event-driven tool execution.
|
||||
* @returns {Record<string, t.EventHandler>} The default handlers.
|
||||
* @throws {Error} If the request is not found.
|
||||
*/
|
||||
|
|
@ -175,6 +132,7 @@ function getDefaultHandlers({
|
|||
toolEndCallback,
|
||||
collectedUsage,
|
||||
streamId = null,
|
||||
toolExecuteOptions = null,
|
||||
}) {
|
||||
if (!res || !aggregateContent) {
|
||||
throw new Error(
|
||||
|
|
@ -184,7 +142,6 @@ function getDefaultHandlers({
|
|||
const handlers = {
|
||||
[GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage),
|
||||
[GraphEvents.TOOL_END]: new ToolEndHandler(toolEndCallback, logger),
|
||||
[GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
|
||||
[GraphEvents.ON_RUN_STEP]: {
|
||||
/**
|
||||
* Handle ON_RUN_STEP event.
|
||||
|
|
@ -192,18 +149,19 @@ function getDefaultHandlers({
|
|||
* @param {StreamEventData} data - The event data.
|
||||
* @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata.
|
||||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
handle: async (event, data, metadata) => {
|
||||
aggregateContent({ event, data });
|
||||
if (data?.stepDetails.type === StepTypes.TOOL_CALLS) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else {
|
||||
const agentName = metadata?.name ?? 'Agent';
|
||||
const isToolCall = data?.stepDetails.type === StepTypes.TOOL_CALLS;
|
||||
const action = isToolCall ? 'performing a task...' : 'thinking...';
|
||||
emitEvent(res, streamId, {
|
||||
await emitEvent(res, streamId, {
|
||||
event: 'on_agent_update',
|
||||
data: {
|
||||
runId: metadata?.run_id,
|
||||
|
|
@ -211,7 +169,6 @@ function getDefaultHandlers({
|
|||
},
|
||||
});
|
||||
}
|
||||
aggregateContent({ event, data });
|
||||
},
|
||||
},
|
||||
[GraphEvents.ON_RUN_STEP_DELTA]: {
|
||||
|
|
@ -221,15 +178,15 @@ function getDefaultHandlers({
|
|||
* @param {StreamEventData} data - The event data.
|
||||
* @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata.
|
||||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (data?.delta.type === StepTypes.TOOL_CALLS) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
handle: async (event, data, metadata) => {
|
||||
aggregateContent({ event, data });
|
||||
if (data?.delta.type === StepTypes.TOOL_CALLS) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
},
|
||||
},
|
||||
[GraphEvents.ON_RUN_STEP_COMPLETED]: {
|
||||
|
|
@ -239,15 +196,15 @@ function getDefaultHandlers({
|
|||
* @param {StreamEventData & { result: ToolEndData }} data - The event data.
|
||||
* @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata.
|
||||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (data?.result != null) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
handle: async (event, data, metadata) => {
|
||||
aggregateContent({ event, data });
|
||||
if (data?.result != null) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
},
|
||||
},
|
||||
[GraphEvents.ON_MESSAGE_DELTA]: {
|
||||
|
|
@ -257,13 +214,13 @@ function getDefaultHandlers({
|
|||
* @param {StreamEventData} data - The event data.
|
||||
* @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata.
|
||||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
handle: async (event, data, metadata) => {
|
||||
aggregateContent({ event, data });
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
},
|
||||
},
|
||||
[GraphEvents.ON_REASONING_DELTA]: {
|
||||
|
|
@ -273,22 +230,27 @@ function getDefaultHandlers({
|
|||
* @param {StreamEventData} data - The event data.
|
||||
* @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata.
|
||||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
handle: async (event, data, metadata) => {
|
||||
aggregateContent({ event, data });
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
if (toolExecuteOptions) {
|
||||
handlers[GraphEvents.ON_TOOL_EXECUTE] = createToolExecuteHandler(toolExecuteOptions);
|
||||
}
|
||||
|
||||
return handlers;
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper to write attachment events either to res or to job emitter.
|
||||
* Note: Attachments are not order-sensitive like deltas, so fire-and-forget is acceptable.
|
||||
* @param {ServerResponse} res - The server response object
|
||||
* @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode
|
||||
* @param {Object} attachment - The attachment data
|
||||
|
|
@ -441,10 +403,10 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null })
|
|||
return;
|
||||
}
|
||||
|
||||
{
|
||||
if (output.name !== Tools.execute_code) {
|
||||
return;
|
||||
}
|
||||
const isCodeTool =
|
||||
output.name === Tools.execute_code || output.name === Constants.PROGRAMMATIC_TOOL_CALLING;
|
||||
if (!isCodeTool) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!output.artifact.files) {
|
||||
|
|
@ -488,7 +450,226 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null })
|
|||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper to write attachment events in Open Responses format (librechat:attachment)
|
||||
* @param {ServerResponse} res - The server response object
|
||||
* @param {Object} tracker - The response tracker with sequence number
|
||||
* @param {Object} attachment - The attachment data
|
||||
* @param {Object} metadata - Additional metadata (messageId, conversationId)
|
||||
*/
|
||||
function writeResponsesAttachment(res, tracker, attachment, metadata) {
|
||||
const sequenceNumber = tracker.nextSequence();
|
||||
writeAttachmentEvent(res, sequenceNumber, attachment, {
|
||||
messageId: metadata.run_id,
|
||||
conversationId: metadata.thread_id,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a tool end callback specifically for the Responses API.
|
||||
* Emits attachments as `librechat:attachment` events per the Open Responses extension spec.
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {ServerResponse} params.res
|
||||
* @param {Object} params.tracker - Response tracker with sequence number
|
||||
* @param {Promise<MongoFile | { filename: string; filepath: string; expires: number;} | null>[]} params.artifactPromises
|
||||
* @returns {ToolEndCallback} The tool end callback.
|
||||
*/
|
||||
function createResponsesToolEndCallback({ req, res, tracker, artifactPromises }) {
|
||||
/**
|
||||
* @type {ToolEndCallback}
|
||||
*/
|
||||
return async (data, metadata) => {
|
||||
const output = data?.output;
|
||||
if (!output) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!output.artifact) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (output.artifact[Tools.file_search]) {
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
const user = req.user;
|
||||
const attachment = await processFileCitations({
|
||||
user,
|
||||
metadata,
|
||||
appConfig: req.config,
|
||||
toolArtifact: output.artifact,
|
||||
toolCallId: output.tool_call_id,
|
||||
});
|
||||
if (!attachment) {
|
||||
return null;
|
||||
}
|
||||
// For Responses API, emit attachment during streaming
|
||||
if (res.headersSent && !res.writableEnded) {
|
||||
writeResponsesAttachment(res, tracker, attachment, metadata);
|
||||
}
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing file citations:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
if (output.artifact[Tools.ui_resources]) {
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
const attachment = {
|
||||
type: Tools.ui_resources,
|
||||
toolCallId: output.tool_call_id,
|
||||
[Tools.ui_resources]: output.artifact[Tools.ui_resources].data,
|
||||
};
|
||||
// For Responses API, always emit attachment during streaming
|
||||
if (res.headersSent && !res.writableEnded) {
|
||||
writeResponsesAttachment(res, tracker, attachment, metadata);
|
||||
}
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
if (output.artifact[Tools.web_search]) {
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
const attachment = {
|
||||
type: Tools.web_search,
|
||||
toolCallId: output.tool_call_id,
|
||||
[Tools.web_search]: { ...output.artifact[Tools.web_search] },
|
||||
};
|
||||
// For Responses API, always emit attachment during streaming
|
||||
if (res.headersSent && !res.writableEnded) {
|
||||
writeResponsesAttachment(res, tracker, attachment, metadata);
|
||||
}
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
if (output.artifact.content) {
|
||||
/** @type {FormattedContent[]} */
|
||||
const content = output.artifact.content;
|
||||
for (let i = 0; i < content.length; i++) {
|
||||
const part = content[i];
|
||||
if (!part) {
|
||||
continue;
|
||||
}
|
||||
if (part.type !== 'image_url') {
|
||||
continue;
|
||||
}
|
||||
const { url } = part.image_url;
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
const filename = `${output.name}_img_${nanoid()}`;
|
||||
const file_id = output.artifact.file_ids?.[i];
|
||||
const file = await saveBase64Image(url, {
|
||||
req,
|
||||
file_id,
|
||||
filename,
|
||||
endpoint: metadata.provider,
|
||||
context: FileContext.image_generation,
|
||||
});
|
||||
const fileMetadata = Object.assign(file, {
|
||||
toolCallId: output.tool_call_id,
|
||||
});
|
||||
|
||||
if (!fileMetadata) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// For Responses API, emit attachment during streaming
|
||||
if (res.headersSent && !res.writableEnded) {
|
||||
const attachment = {
|
||||
file_id: fileMetadata.file_id,
|
||||
filename: fileMetadata.filename,
|
||||
type: fileMetadata.type,
|
||||
url: fileMetadata.filepath,
|
||||
width: fileMetadata.width,
|
||||
height: fileMetadata.height,
|
||||
tool_call_id: output.tool_call_id,
|
||||
};
|
||||
writeResponsesAttachment(res, tracker, attachment, metadata);
|
||||
}
|
||||
|
||||
return fileMetadata;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const isCodeTool =
|
||||
output.name === Tools.execute_code || output.name === Constants.PROGRAMMATIC_TOOL_CALLING;
|
||||
if (!isCodeTool) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!output.artifact.files) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (const file of output.artifact.files) {
|
||||
const { id, name } = file;
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
const result = await loadAuthValues({
|
||||
userId: req.user.id,
|
||||
authFields: [EnvVar.CODE_API_KEY],
|
||||
});
|
||||
const fileMetadata = await processCodeOutput({
|
||||
req,
|
||||
id,
|
||||
name,
|
||||
apiKey: result[EnvVar.CODE_API_KEY],
|
||||
messageId: metadata.run_id,
|
||||
toolCallId: output.tool_call_id,
|
||||
conversationId: metadata.thread_id,
|
||||
session_id: output.artifact.session_id,
|
||||
});
|
||||
|
||||
if (!fileMetadata) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// For Responses API, emit attachment during streaming
|
||||
if (res.headersSent && !res.writableEnded) {
|
||||
const attachment = {
|
||||
file_id: fileMetadata.file_id,
|
||||
filename: fileMetadata.filename,
|
||||
type: fileMetadata.type,
|
||||
url: fileMetadata.filepath,
|
||||
width: fileMetadata.width,
|
||||
height: fileMetadata.height,
|
||||
tool_call_id: output.tool_call_id,
|
||||
};
|
||||
writeResponsesAttachment(res, tracker, attachment, metadata);
|
||||
}
|
||||
|
||||
return fileMetadata;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing code output:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getDefaultHandlers,
|
||||
createToolEndCallback,
|
||||
createResponsesToolEndCallback,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,22 +1,27 @@
|
|||
require('events').EventEmitter.defaultMaxListeners = 100;
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { DynamicStructuredTool } = require('@langchain/core/tools');
|
||||
const { getBufferString, HumanMessage } = require('@langchain/core/messages');
|
||||
const {
|
||||
createRun,
|
||||
Tokenizer,
|
||||
checkAccess,
|
||||
logAxiosError,
|
||||
buildToolSet,
|
||||
sanitizeTitle,
|
||||
logToolError,
|
||||
payloadParser,
|
||||
resolveHeaders,
|
||||
createSafeUser,
|
||||
initializeAgent,
|
||||
getBalanceConfig,
|
||||
getProviderConfig,
|
||||
omitTitleOptions,
|
||||
memoryInstructions,
|
||||
applyContextToAgent,
|
||||
createTokenCounter,
|
||||
GenerationJobManager,
|
||||
getTransactionsConfig,
|
||||
createMemoryProcessor,
|
||||
createMultiAgentMapper,
|
||||
filterMalformedContentParts,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
|
|
@ -24,9 +29,7 @@ const {
|
|||
Providers,
|
||||
TitleMethod,
|
||||
formatMessage,
|
||||
labelContentByAgent,
|
||||
formatAgentMessages,
|
||||
getTokenCountForMessage,
|
||||
createMetadataAggregator,
|
||||
} = require('@librechat/agents');
|
||||
const {
|
||||
|
|
@ -38,7 +41,6 @@ const {
|
|||
PermissionTypes,
|
||||
isAgentsEndpoint,
|
||||
isEphemeralAgentId,
|
||||
bedrockInputSchema,
|
||||
removeNullishValues,
|
||||
} = require('librechat-data-provider');
|
||||
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
|
||||
|
|
@ -51,183 +53,6 @@ const { loadAgent } = require('~/models/Agent');
|
|||
const { getMCPManager } = require('~/config');
|
||||
const db = require('~/models');
|
||||
|
||||
const omitTitleOptions = new Set([
|
||||
'stream',
|
||||
'thinking',
|
||||
'streaming',
|
||||
'clientOptions',
|
||||
'thinkingConfig',
|
||||
'thinkingBudget',
|
||||
'includeThoughts',
|
||||
'maxOutputTokens',
|
||||
'additionalModelRequestFields',
|
||||
]);
|
||||
|
||||
/**
|
||||
* @param {ServerRequest} req
|
||||
* @param {Agent} agent
|
||||
* @param {string} endpoint
|
||||
*/
|
||||
const payloadParser = ({ req, agent, endpoint }) => {
|
||||
if (isAgentsEndpoint(endpoint)) {
|
||||
return { model: undefined };
|
||||
} else if (endpoint === EModelEndpoint.bedrock) {
|
||||
const parsedValues = bedrockInputSchema.parse(agent.model_parameters);
|
||||
if (parsedValues.thinking == null) {
|
||||
parsedValues.thinking = false;
|
||||
}
|
||||
return parsedValues;
|
||||
}
|
||||
return req.body.endpointOption.model_parameters;
|
||||
};
|
||||
|
||||
function createTokenCounter(encoding) {
|
||||
return function (message) {
|
||||
const countTokens = (text) => Tokenizer.getTokenCount(text, encoding);
|
||||
return getTokenCountForMessage(message, countTokens);
|
||||
};
|
||||
}
|
||||
|
||||
function logToolError(graph, error, toolId) {
|
||||
logAxiosError({
|
||||
error,
|
||||
message: `[api/server/controllers/agents/client.js #chatCompletion] Tool Error "${toolId}"`,
|
||||
});
|
||||
}
|
||||
|
||||
/** Regex pattern to match agent ID suffix (____N) */
|
||||
const AGENT_SUFFIX_PATTERN = /____(\d+)$/;
|
||||
|
||||
/**
|
||||
* Finds the primary agent ID within a set of agent IDs.
|
||||
* Primary = no suffix (____N) or lowest suffix number.
|
||||
* @param {Set<string>} agentIds
|
||||
* @returns {string | null}
|
||||
*/
|
||||
function findPrimaryAgentId(agentIds) {
|
||||
let primaryAgentId = null;
|
||||
let lowestSuffixIndex = Infinity;
|
||||
|
||||
for (const agentId of agentIds) {
|
||||
const suffixMatch = agentId.match(AGENT_SUFFIX_PATTERN);
|
||||
if (!suffixMatch) {
|
||||
return agentId;
|
||||
}
|
||||
const suffixIndex = parseInt(suffixMatch[1], 10);
|
||||
if (suffixIndex < lowestSuffixIndex) {
|
||||
lowestSuffixIndex = suffixIndex;
|
||||
primaryAgentId = agentId;
|
||||
}
|
||||
}
|
||||
|
||||
return primaryAgentId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a mapMethod for getMessagesForConversation that processes agent content.
|
||||
* - Strips agentId/groupId metadata from all content
|
||||
* - For parallel agents (addedConvo with groupId): filters each group to its primary agent
|
||||
* - For handoffs (agentId without groupId): keeps all content from all agents
|
||||
* - For multi-agent: applies agent labels to content
|
||||
*
|
||||
* The key distinction:
|
||||
* - Parallel execution (addedConvo): Parts have both agentId AND groupId
|
||||
* - Handoffs: Parts only have agentId, no groupId
|
||||
*
|
||||
* @param {Agent} primaryAgent - Primary agent configuration
|
||||
* @param {Map<string, Agent>} [agentConfigs] - Additional agent configurations
|
||||
* @returns {(message: TMessage) => TMessage} Map method for processing messages
|
||||
*/
|
||||
function createMultiAgentMapper(primaryAgent, agentConfigs) {
|
||||
const hasMultipleAgents = (primaryAgent.edges?.length ?? 0) > 0 || (agentConfigs?.size ?? 0) > 0;
|
||||
|
||||
/** @type {Record<string, string> | null} */
|
||||
let agentNames = null;
|
||||
if (hasMultipleAgents) {
|
||||
agentNames = { [primaryAgent.id]: primaryAgent.name || 'Assistant' };
|
||||
if (agentConfigs) {
|
||||
for (const [agentId, agentConfig] of agentConfigs.entries()) {
|
||||
agentNames[agentId] = agentConfig.name || agentConfig.id;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return (message) => {
|
||||
if (message.isCreatedByUser || !Array.isArray(message.content)) {
|
||||
return message;
|
||||
}
|
||||
|
||||
// Check for metadata
|
||||
const hasAgentMetadata = message.content.some((part) => part?.agentId || part?.groupId != null);
|
||||
if (!hasAgentMetadata) {
|
||||
return message;
|
||||
}
|
||||
|
||||
try {
|
||||
// Build a map of groupId -> Set of agentIds, to find primary per group
|
||||
/** @type {Map<number, Set<string>>} */
|
||||
const groupAgentMap = new Map();
|
||||
|
||||
for (const part of message.content) {
|
||||
const groupId = part?.groupId;
|
||||
const agentId = part?.agentId;
|
||||
if (groupId != null && agentId) {
|
||||
if (!groupAgentMap.has(groupId)) {
|
||||
groupAgentMap.set(groupId, new Set());
|
||||
}
|
||||
groupAgentMap.get(groupId).add(agentId);
|
||||
}
|
||||
}
|
||||
|
||||
// For each group, find the primary agent
|
||||
/** @type {Map<number, string>} */
|
||||
const groupPrimaryMap = new Map();
|
||||
for (const [groupId, agentIds] of groupAgentMap) {
|
||||
const primary = findPrimaryAgentId(agentIds);
|
||||
if (primary) {
|
||||
groupPrimaryMap.set(groupId, primary);
|
||||
}
|
||||
}
|
||||
|
||||
/** @type {Array<TMessageContentParts>} */
|
||||
const filteredContent = [];
|
||||
/** @type {Record<number, string>} */
|
||||
const agentIdMap = {};
|
||||
|
||||
for (const part of message.content) {
|
||||
const agentId = part?.agentId;
|
||||
const groupId = part?.groupId;
|
||||
|
||||
// Filtering logic:
|
||||
// - No groupId (handoffs): always include
|
||||
// - Has groupId (parallel): only include if it's the primary for that group
|
||||
const isParallelPart = groupId != null;
|
||||
const groupPrimary = isParallelPart ? groupPrimaryMap.get(groupId) : null;
|
||||
const shouldInclude = !isParallelPart || !agentId || agentId === groupPrimary;
|
||||
|
||||
if (shouldInclude) {
|
||||
const newIndex = filteredContent.length;
|
||||
const { agentId: _a, groupId: _g, ...cleanPart } = part;
|
||||
filteredContent.push(cleanPart);
|
||||
if (agentId && hasMultipleAgents) {
|
||||
agentIdMap[newIndex] = agentId;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const finalContent =
|
||||
Object.keys(agentIdMap).length > 0 && agentNames
|
||||
? labelContentByAgent(filteredContent, agentIdMap, agentNames)
|
||||
: filteredContent;
|
||||
|
||||
return { ...message, content: finalContent };
|
||||
} catch (error) {
|
||||
logger.error('[AgentClient] Error processing multi-agent message:', error);
|
||||
return message;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
class AgentClient extends BaseClient {
|
||||
constructor(options = {}) {
|
||||
super(null, options);
|
||||
|
|
@ -295,14 +120,9 @@ class AgentClient extends BaseClient {
|
|||
checkVisionRequest() {}
|
||||
|
||||
getSaveOptions() {
|
||||
// TODO:
|
||||
// would need to be override settings; otherwise, model needs to be undefined
|
||||
// model: this.override.model,
|
||||
// instructions: this.override.instructions,
|
||||
// additional_instructions: this.override.additional_instructions,
|
||||
let runOptions = {};
|
||||
try {
|
||||
runOptions = payloadParser(this.options);
|
||||
runOptions = payloadParser(this.options) ?? {};
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
'[api/server/controllers/agents/client.js #getSaveOptions] Error parsing options',
|
||||
|
|
@ -313,14 +133,14 @@ class AgentClient extends BaseClient {
|
|||
return removeNullishValues(
|
||||
Object.assign(
|
||||
{
|
||||
spec: this.options.spec,
|
||||
iconURL: this.options.iconURL,
|
||||
endpoint: this.options.endpoint,
|
||||
agent_id: this.options.agent.id,
|
||||
modelLabel: this.options.modelLabel,
|
||||
maxContextTokens: this.options.maxContextTokens,
|
||||
resendFiles: this.options.resendFiles,
|
||||
imageDetail: this.options.imageDetail,
|
||||
spec: this.options.spec,
|
||||
iconURL: this.options.iconURL,
|
||||
maxContextTokens: this.maxContextTokens,
|
||||
},
|
||||
// TODO: PARSE OPTIONS BY PROVIDER, MAY CONTAIN SENSITIVE DATA
|
||||
runOptions,
|
||||
|
|
@ -328,11 +148,13 @@ class AgentClient extends BaseClient {
|
|||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns build message options. For AgentClient, agent-specific instructions
|
||||
* are retrieved directly from agent objects in buildMessages, so this returns empty.
|
||||
* @returns {Object} Empty options object
|
||||
*/
|
||||
getBuildMessagesOptions() {
|
||||
return {
|
||||
instructions: this.options.agent.instructions,
|
||||
additional_instructions: this.options.agent.additional_instructions,
|
||||
};
|
||||
return {};
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -355,12 +177,7 @@ class AgentClient extends BaseClient {
|
|||
return files;
|
||||
}
|
||||
|
||||
async buildMessages(
|
||||
messages,
|
||||
parentMessageId,
|
||||
{ instructions = null, additional_instructions = null },
|
||||
opts,
|
||||
) {
|
||||
async buildMessages(messages, parentMessageId, _buildOptions, opts) {
|
||||
/** Always pass mapMethod; getMessagesForConversation applies it only to messages with addedConvo flag */
|
||||
const orderedMessages = this.constructor.getMessagesForConversation({
|
||||
messages,
|
||||
|
|
@ -374,11 +191,29 @@ class AgentClient extends BaseClient {
|
|||
/** @type {number | undefined} */
|
||||
let promptTokens;
|
||||
|
||||
/** @type {string} */
|
||||
let systemContent = [instructions ?? '', additional_instructions ?? '']
|
||||
.filter(Boolean)
|
||||
.join('\n')
|
||||
.trim();
|
||||
/**
|
||||
* Extract base instructions for all agents (combines instructions + additional_instructions).
|
||||
* This must be done before applying context to preserve the original agent configuration.
|
||||
*/
|
||||
const extractBaseInstructions = (agent) => {
|
||||
const baseInstructions = [agent.instructions ?? '', agent.additional_instructions ?? '']
|
||||
.filter(Boolean)
|
||||
.join('\n')
|
||||
.trim();
|
||||
agent.instructions = baseInstructions;
|
||||
return agent;
|
||||
};
|
||||
|
||||
/** Collect all agents for unified processing, extracting base instructions during collection */
|
||||
const allAgents = [
|
||||
{ agent: extractBaseInstructions(this.options.agent), agentId: this.options.agent.id },
|
||||
...(this.agentConfigs?.size > 0
|
||||
? Array.from(this.agentConfigs.entries()).map(([agentId, agent]) => ({
|
||||
agent: extractBaseInstructions(agent),
|
||||
agentId,
|
||||
}))
|
||||
: []),
|
||||
];
|
||||
|
||||
if (this.options.attachments) {
|
||||
const attachments = await this.options.attachments;
|
||||
|
|
@ -413,6 +248,7 @@ class AgentClient extends BaseClient {
|
|||
assistantName: this.options?.modelLabel,
|
||||
});
|
||||
|
||||
/** For non-latest messages, prepend file context directly to message content */
|
||||
if (message.fileContext && i !== orderedMessages.length - 1) {
|
||||
if (typeof formattedMessage.content === 'string') {
|
||||
formattedMessage.content = message.fileContext + '\n' + formattedMessage.content;
|
||||
|
|
@ -422,8 +258,6 @@ class AgentClient extends BaseClient {
|
|||
? (textPart.text = message.fileContext + '\n' + textPart.text)
|
||||
: formattedMessage.content.unshift({ type: 'text', text: message.fileContext });
|
||||
}
|
||||
} else if (message.fileContext && i === orderedMessages.length - 1) {
|
||||
systemContent = [systemContent, message.fileContext].join('\n');
|
||||
}
|
||||
|
||||
const needsTokenCount =
|
||||
|
|
@ -456,46 +290,35 @@ class AgentClient extends BaseClient {
|
|||
return formattedMessage;
|
||||
});
|
||||
|
||||
/**
|
||||
* Build shared run context - applies to ALL agents in the run.
|
||||
* This includes: file context (latest message), augmented prompt (RAG), memory context.
|
||||
*/
|
||||
const sharedRunContextParts = [];
|
||||
|
||||
/** File context from the latest message (attachments) */
|
||||
const latestMessage = orderedMessages[orderedMessages.length - 1];
|
||||
if (latestMessage?.fileContext) {
|
||||
sharedRunContextParts.push(latestMessage.fileContext);
|
||||
}
|
||||
|
||||
/** Augmented prompt from RAG/context handlers */
|
||||
if (this.contextHandlers) {
|
||||
this.augmentedPrompt = await this.contextHandlers.createContext();
|
||||
systemContent = this.augmentedPrompt + systemContent;
|
||||
}
|
||||
|
||||
// Inject MCP server instructions if available
|
||||
const ephemeralAgent = this.options.req.body.ephemeralAgent;
|
||||
let mcpServers = [];
|
||||
|
||||
// Check for ephemeral agent MCP servers
|
||||
if (ephemeralAgent && ephemeralAgent.mcp && ephemeralAgent.mcp.length > 0) {
|
||||
mcpServers = ephemeralAgent.mcp;
|
||||
}
|
||||
// Check for regular agent MCP tools
|
||||
else if (this.options.agent && this.options.agent.tools) {
|
||||
mcpServers = this.options.agent.tools
|
||||
.filter(
|
||||
(tool) =>
|
||||
tool instanceof DynamicStructuredTool && tool.name.includes(Constants.mcp_delimiter),
|
||||
)
|
||||
.map((tool) => tool.name.split(Constants.mcp_delimiter).pop())
|
||||
.filter(Boolean);
|
||||
}
|
||||
|
||||
if (mcpServers.length > 0) {
|
||||
try {
|
||||
const mcpInstructions = await getMCPManager().formatInstructionsForContext(mcpServers);
|
||||
if (mcpInstructions) {
|
||||
systemContent = [systemContent, mcpInstructions].filter(Boolean).join('\n\n');
|
||||
logger.debug('[AgentClient] Injected MCP instructions for servers:', mcpServers);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('[AgentClient] Failed to inject MCP instructions:', error);
|
||||
if (this.augmentedPrompt) {
|
||||
sharedRunContextParts.push(this.augmentedPrompt);
|
||||
}
|
||||
}
|
||||
|
||||
if (systemContent) {
|
||||
this.options.agent.instructions = systemContent;
|
||||
/** Memory context (user preferences/memories) */
|
||||
const withoutKeys = await this.useMemory();
|
||||
if (withoutKeys) {
|
||||
const memoryContext = `${memoryInstructions}\n\n# Existing memory about the user:\n${withoutKeys}`;
|
||||
sharedRunContextParts.push(memoryContext);
|
||||
}
|
||||
|
||||
const sharedRunContext = sharedRunContextParts.join('\n\n');
|
||||
|
||||
/** @type {Record<string, number> | undefined} */
|
||||
let tokenCountMap;
|
||||
|
||||
|
|
@ -521,14 +344,27 @@ class AgentClient extends BaseClient {
|
|||
opts.getReqData({ promptTokens });
|
||||
}
|
||||
|
||||
const withoutKeys = await this.useMemory();
|
||||
if (withoutKeys) {
|
||||
systemContent += `${memoryInstructions}\n\n# Existing memory about the user:\n${withoutKeys}`;
|
||||
}
|
||||
|
||||
if (systemContent) {
|
||||
this.options.agent.instructions = systemContent;
|
||||
}
|
||||
/**
|
||||
* Apply context to all agents.
|
||||
* Each agent gets: shared run context + their own base instructions + their own MCP instructions.
|
||||
*
|
||||
* NOTE: This intentionally mutates agent objects in place. The agentConfigs Map
|
||||
* holds references to config objects that will be passed to the graph runtime.
|
||||
*/
|
||||
const ephemeralAgent = this.options.req.body.ephemeralAgent;
|
||||
const mcpManager = getMCPManager();
|
||||
await Promise.all(
|
||||
allAgents.map(({ agent, agentId }) =>
|
||||
applyContextToAgent({
|
||||
agent,
|
||||
agentId,
|
||||
logger,
|
||||
mcpManager,
|
||||
sharedRunContext,
|
||||
ephemeralAgent: agentId === this.options.agent.id ? ephemeralAgent : undefined,
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
|
@ -600,6 +436,8 @@ class AgentClient extends BaseClient {
|
|||
agent_id: memoryConfig.agent.id,
|
||||
endpoint: EModelEndpoint.agents,
|
||||
});
|
||||
} else if (memoryConfig.agent?.id != null) {
|
||||
prelimAgent = this.options.agent;
|
||||
} else if (
|
||||
memoryConfig.agent?.id == null &&
|
||||
memoryConfig.agent?.model != null &&
|
||||
|
|
@ -614,6 +452,10 @@ class AgentClient extends BaseClient {
|
|||
);
|
||||
}
|
||||
|
||||
if (!prelimAgent) {
|
||||
return;
|
||||
}
|
||||
|
||||
const agent = await initializeAgent(
|
||||
{
|
||||
req: this.options.req,
|
||||
|
|
@ -633,6 +475,7 @@ class AgentClient extends BaseClient {
|
|||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -945,13 +788,13 @@ class AgentClient extends BaseClient {
|
|||
},
|
||||
user: createSafeUser(this.options.req.user),
|
||||
},
|
||||
recursionLimit: agentsEConfig?.recursionLimit ?? 25,
|
||||
recursionLimit: agentsEConfig?.recursionLimit ?? 50,
|
||||
signal: abortController.signal,
|
||||
streamMode: 'values',
|
||||
version: 'v2',
|
||||
};
|
||||
|
||||
const toolSet = new Set((this.options.agent.tools ?? []).map((tool) => tool && tool.name));
|
||||
const toolSet = buildToolSet(this.options.agent);
|
||||
let { messages: initialMessages, indexTokenCountMap } = formatAgentMessages(
|
||||
payload,
|
||||
this.indexTokenCountMap,
|
||||
|
|
@ -1012,6 +855,7 @@ class AgentClient extends BaseClient {
|
|||
|
||||
run = await createRun({
|
||||
agents,
|
||||
messages,
|
||||
indexTokenCountMap,
|
||||
runId: this.responseMessageId,
|
||||
signal: abortController.signal,
|
||||
|
|
@ -1084,11 +928,20 @@ class AgentClient extends BaseClient {
|
|||
this.artifactPromises.push(...attachments);
|
||||
}
|
||||
|
||||
await this.recordCollectedUsage({
|
||||
context: 'message',
|
||||
balance: balanceConfig,
|
||||
transactions: transactionsConfig,
|
||||
});
|
||||
/** Skip token spending if aborted - the abort handler (abortMiddleware.js) handles it
|
||||
This prevents double-spending when user aborts via `/api/agents/chat/abort` */
|
||||
const wasAborted = abortController?.signal?.aborted;
|
||||
if (!wasAborted) {
|
||||
await this.recordCollectedUsage({
|
||||
context: 'message',
|
||||
balance: balanceConfig,
|
||||
transactions: transactionsConfig,
|
||||
});
|
||||
} else {
|
||||
logger.debug(
|
||||
'[api/server/controllers/agents/client.js #chatCompletion] Skipping token spending - handled by abort middleware',
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error(
|
||||
'[api/server/controllers/agents/client.js #chatCompletion] Error in cleanup phase',
|
||||
|
|
|
|||
|
|
@ -12,6 +12,17 @@ jest.mock('@librechat/agents', () => ({
|
|||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
...jest.requireActual('@librechat/api'),
|
||||
checkAccess: jest.fn(),
|
||||
initializeAgent: jest.fn(),
|
||||
createMemoryProcessor: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Agent', () => ({
|
||||
loadAgent: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/models/Role', () => ({
|
||||
getRoleByName: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock getMCPManager
|
||||
|
|
@ -1310,8 +1321,8 @@ describe('AgentClient - titleConvo', () => {
|
|||
expect(client.options.agent.instructions).toContain('# MCP Server Instructions');
|
||||
expect(client.options.agent.instructions).toContain('Use these tools carefully');
|
||||
|
||||
// Verify the base instructions are also included
|
||||
expect(client.options.agent.instructions).toContain('Base instructions');
|
||||
// Verify the base instructions are also included (from agent config, not buildOptions)
|
||||
expect(client.options.agent.instructions).toContain('Base agent instructions');
|
||||
});
|
||||
|
||||
it('should handle MCP instructions with ephemeral agent', async () => {
|
||||
|
|
@ -1373,8 +1384,8 @@ describe('AgentClient - titleConvo', () => {
|
|||
additional_instructions: null,
|
||||
});
|
||||
|
||||
// Verify the instructions still work without MCP content
|
||||
expect(client.options.agent.instructions).toBe('Base instructions only');
|
||||
// Verify the instructions still work without MCP content (from agent config, not buildOptions)
|
||||
expect(client.options.agent.instructions).toBe('Base agent instructions');
|
||||
expect(client.options.agent.instructions).not.toContain('[object Promise]');
|
||||
});
|
||||
|
||||
|
|
@ -1398,8 +1409,8 @@ describe('AgentClient - titleConvo', () => {
|
|||
additional_instructions: null,
|
||||
});
|
||||
|
||||
// Should still have base instructions without MCP content
|
||||
expect(client.options.agent.instructions).toContain('Base instructions');
|
||||
// Should still have base instructions without MCP content (from agent config, not buildOptions)
|
||||
expect(client.options.agent.instructions).toContain('Base agent instructions');
|
||||
expect(client.options.agent.instructions).not.toContain('[object Promise]');
|
||||
});
|
||||
});
|
||||
|
|
@ -1849,4 +1860,400 @@ describe('AgentClient - titleConvo', () => {
|
|||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('buildMessages - memory context for parallel agents', () => {
|
||||
let client;
|
||||
let mockReq;
|
||||
let mockRes;
|
||||
let mockAgent;
|
||||
let mockOptions;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
|
||||
mockAgent = {
|
||||
id: 'primary-agent',
|
||||
name: 'Primary Agent',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
provider: EModelEndpoint.openAI,
|
||||
instructions: 'Primary agent instructions',
|
||||
model_parameters: {
|
||||
model: 'gpt-4',
|
||||
},
|
||||
tools: [],
|
||||
};
|
||||
|
||||
mockReq = {
|
||||
user: {
|
||||
id: 'user-123',
|
||||
personalization: {
|
||||
memories: true,
|
||||
},
|
||||
},
|
||||
body: {
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
},
|
||||
config: {
|
||||
memory: {
|
||||
disabled: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
mockRes = {};
|
||||
|
||||
mockOptions = {
|
||||
req: mockReq,
|
||||
res: mockRes,
|
||||
agent: mockAgent,
|
||||
endpoint: EModelEndpoint.agents,
|
||||
};
|
||||
|
||||
client = new AgentClient(mockOptions);
|
||||
client.conversationId = 'convo-123';
|
||||
client.responseMessageId = 'response-123';
|
||||
client.shouldSummarize = false;
|
||||
client.maxContextTokens = 4096;
|
||||
});
|
||||
|
||||
it('should pass memory context to parallel agents (addedConvo)', async () => {
|
||||
const memoryContent = 'User prefers dark mode. User is a software developer.';
|
||||
client.useMemory = jest.fn().mockResolvedValue(memoryContent);
|
||||
|
||||
const parallelAgent1 = {
|
||||
id: 'parallel-agent-1',
|
||||
name: 'Parallel Agent 1',
|
||||
instructions: 'Parallel agent 1 instructions',
|
||||
provider: EModelEndpoint.openAI,
|
||||
};
|
||||
|
||||
const parallelAgent2 = {
|
||||
id: 'parallel-agent-2',
|
||||
name: 'Parallel Agent 2',
|
||||
instructions: 'Parallel agent 2 instructions',
|
||||
provider: EModelEndpoint.anthropic,
|
||||
};
|
||||
|
||||
client.agentConfigs = new Map([
|
||||
['parallel-agent-1', parallelAgent1],
|
||||
['parallel-agent-2', parallelAgent2],
|
||||
]);
|
||||
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
];
|
||||
|
||||
await client.buildMessages(messages, null, {
|
||||
instructions: 'Base instructions',
|
||||
additional_instructions: null,
|
||||
});
|
||||
|
||||
expect(client.useMemory).toHaveBeenCalled();
|
||||
|
||||
// Verify primary agent has its configured instructions (not from buildOptions) and memory context
|
||||
expect(client.options.agent.instructions).toContain('Primary agent instructions');
|
||||
expect(client.options.agent.instructions).toContain(memoryContent);
|
||||
|
||||
expect(parallelAgent1.instructions).toContain('Parallel agent 1 instructions');
|
||||
expect(parallelAgent1.instructions).toContain(memoryContent);
|
||||
|
||||
expect(parallelAgent2.instructions).toContain('Parallel agent 2 instructions');
|
||||
expect(parallelAgent2.instructions).toContain(memoryContent);
|
||||
});
|
||||
|
||||
it('should not modify parallel agents when no memory context is available', async () => {
|
||||
client.useMemory = jest.fn().mockResolvedValue(undefined);
|
||||
|
||||
const parallelAgent = {
|
||||
id: 'parallel-agent-1',
|
||||
name: 'Parallel Agent 1',
|
||||
instructions: 'Original parallel instructions',
|
||||
provider: EModelEndpoint.openAI,
|
||||
};
|
||||
|
||||
client.agentConfigs = new Map([['parallel-agent-1', parallelAgent]]);
|
||||
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
];
|
||||
|
||||
await client.buildMessages(messages, null, {
|
||||
instructions: 'Base instructions',
|
||||
additional_instructions: null,
|
||||
});
|
||||
|
||||
expect(parallelAgent.instructions).toBe('Original parallel instructions');
|
||||
});
|
||||
|
||||
it('should handle parallel agents without existing instructions', async () => {
|
||||
const memoryContent = 'User is a data scientist.';
|
||||
client.useMemory = jest.fn().mockResolvedValue(memoryContent);
|
||||
|
||||
const parallelAgentNoInstructions = {
|
||||
id: 'parallel-agent-no-instructions',
|
||||
name: 'Parallel Agent No Instructions',
|
||||
provider: EModelEndpoint.openAI,
|
||||
};
|
||||
|
||||
client.agentConfigs = new Map([
|
||||
['parallel-agent-no-instructions', parallelAgentNoInstructions],
|
||||
]);
|
||||
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
];
|
||||
|
||||
await client.buildMessages(messages, null, {
|
||||
instructions: null,
|
||||
additional_instructions: null,
|
||||
});
|
||||
|
||||
expect(parallelAgentNoInstructions.instructions).toContain(memoryContent);
|
||||
});
|
||||
|
||||
it('should not modify agentConfigs when none exist', async () => {
|
||||
const memoryContent = 'User prefers concise responses.';
|
||||
client.useMemory = jest.fn().mockResolvedValue(memoryContent);
|
||||
|
||||
client.agentConfigs = null;
|
||||
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
];
|
||||
|
||||
await expect(
|
||||
client.buildMessages(messages, null, {
|
||||
instructions: 'Base instructions',
|
||||
additional_instructions: null,
|
||||
}),
|
||||
).resolves.not.toThrow();
|
||||
|
||||
expect(client.options.agent.instructions).toContain(memoryContent);
|
||||
});
|
||||
|
||||
it('should handle empty agentConfigs map', async () => {
|
||||
const memoryContent = 'User likes detailed explanations.';
|
||||
client.useMemory = jest.fn().mockResolvedValue(memoryContent);
|
||||
|
||||
client.agentConfigs = new Map();
|
||||
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
];
|
||||
|
||||
await expect(
|
||||
client.buildMessages(messages, null, {
|
||||
instructions: 'Base instructions',
|
||||
additional_instructions: null,
|
||||
}),
|
||||
).resolves.not.toThrow();
|
||||
|
||||
expect(client.options.agent.instructions).toContain(memoryContent);
|
||||
});
|
||||
});
|
||||
|
||||
describe('useMemory method - prelimAgent assignment', () => {
|
||||
let client;
|
||||
let mockReq;
|
||||
let mockRes;
|
||||
let mockAgent;
|
||||
let mockOptions;
|
||||
let mockCheckAccess;
|
||||
let mockLoadAgent;
|
||||
let mockInitializeAgent;
|
||||
let mockCreateMemoryProcessor;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
|
||||
mockAgent = {
|
||||
id: 'agent-123',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
provider: EModelEndpoint.openAI,
|
||||
instructions: 'Test instructions',
|
||||
model: 'gpt-4',
|
||||
model_parameters: {
|
||||
model: 'gpt-4',
|
||||
},
|
||||
};
|
||||
|
||||
mockReq = {
|
||||
user: {
|
||||
id: 'user-123',
|
||||
personalization: {
|
||||
memories: true,
|
||||
},
|
||||
},
|
||||
config: {
|
||||
memory: {
|
||||
agent: {
|
||||
id: 'agent-123',
|
||||
},
|
||||
},
|
||||
endpoints: {
|
||||
[EModelEndpoint.agents]: {
|
||||
allowedProviders: [EModelEndpoint.openAI],
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
mockRes = {};
|
||||
|
||||
mockOptions = {
|
||||
req: mockReq,
|
||||
res: mockRes,
|
||||
agent: mockAgent,
|
||||
};
|
||||
|
||||
mockCheckAccess = require('@librechat/api').checkAccess;
|
||||
mockLoadAgent = require('~/models/Agent').loadAgent;
|
||||
mockInitializeAgent = require('@librechat/api').initializeAgent;
|
||||
mockCreateMemoryProcessor = require('@librechat/api').createMemoryProcessor;
|
||||
});
|
||||
|
||||
it('should use current agent when memory config agent.id matches current agent id', async () => {
|
||||
mockCheckAccess.mockResolvedValue(true);
|
||||
mockInitializeAgent.mockResolvedValue({
|
||||
...mockAgent,
|
||||
provider: EModelEndpoint.openAI,
|
||||
});
|
||||
mockCreateMemoryProcessor.mockResolvedValue([undefined, jest.fn()]);
|
||||
|
||||
client = new AgentClient(mockOptions);
|
||||
client.conversationId = 'convo-123';
|
||||
client.responseMessageId = 'response-123';
|
||||
|
||||
await client.useMemory();
|
||||
|
||||
expect(mockLoadAgent).not.toHaveBeenCalled();
|
||||
expect(mockInitializeAgent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent: mockAgent,
|
||||
}),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it('should load different agent when memory config agent.id differs from current agent id', async () => {
|
||||
const differentAgentId = 'different-agent-456';
|
||||
const differentAgent = {
|
||||
id: differentAgentId,
|
||||
provider: EModelEndpoint.openAI,
|
||||
model: 'gpt-4',
|
||||
instructions: 'Different agent instructions',
|
||||
};
|
||||
|
||||
mockReq.config.memory.agent.id = differentAgentId;
|
||||
|
||||
mockCheckAccess.mockResolvedValue(true);
|
||||
mockLoadAgent.mockResolvedValue(differentAgent);
|
||||
mockInitializeAgent.mockResolvedValue({
|
||||
...differentAgent,
|
||||
provider: EModelEndpoint.openAI,
|
||||
});
|
||||
mockCreateMemoryProcessor.mockResolvedValue([undefined, jest.fn()]);
|
||||
|
||||
client = new AgentClient(mockOptions);
|
||||
client.conversationId = 'convo-123';
|
||||
client.responseMessageId = 'response-123';
|
||||
|
||||
await client.useMemory();
|
||||
|
||||
expect(mockLoadAgent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent_id: differentAgentId,
|
||||
}),
|
||||
);
|
||||
expect(mockInitializeAgent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent: differentAgent,
|
||||
}),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it('should return early when prelimAgent is undefined (no valid memory agent config)', async () => {
|
||||
mockReq.config.memory = {
|
||||
agent: {},
|
||||
};
|
||||
|
||||
mockCheckAccess.mockResolvedValue(true);
|
||||
|
||||
client = new AgentClient(mockOptions);
|
||||
client.conversationId = 'convo-123';
|
||||
client.responseMessageId = 'response-123';
|
||||
|
||||
const result = await client.useMemory();
|
||||
|
||||
expect(result).toBeUndefined();
|
||||
expect(mockInitializeAgent).not.toHaveBeenCalled();
|
||||
expect(mockCreateMemoryProcessor).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should create ephemeral agent when no id but model and provider are specified', async () => {
|
||||
mockReq.config.memory = {
|
||||
agent: {
|
||||
model: 'gpt-4',
|
||||
provider: EModelEndpoint.openAI,
|
||||
},
|
||||
};
|
||||
|
||||
mockCheckAccess.mockResolvedValue(true);
|
||||
mockInitializeAgent.mockResolvedValue({
|
||||
id: Constants.EPHEMERAL_AGENT_ID,
|
||||
model: 'gpt-4',
|
||||
provider: EModelEndpoint.openAI,
|
||||
});
|
||||
mockCreateMemoryProcessor.mockResolvedValue([undefined, jest.fn()]);
|
||||
|
||||
client = new AgentClient(mockOptions);
|
||||
client.conversationId = 'convo-123';
|
||||
client.responseMessageId = 'response-123';
|
||||
|
||||
await client.useMemory();
|
||||
|
||||
expect(mockLoadAgent).not.toHaveBeenCalled();
|
||||
expect(mockInitializeAgent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent: expect.objectContaining({
|
||||
id: Constants.EPHEMERAL_AGENT_ID,
|
||||
model: 'gpt-4',
|
||||
provider: EModelEndpoint.openAI,
|
||||
}),
|
||||
}),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
701
api/server/controllers/agents/openai.js
Normal file
701
api/server/controllers/agents/openai.js
Normal file
|
|
@ -0,0 +1,701 @@
|
|||
const { nanoid } = require('nanoid');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Callback, ToolEndHandler, formatAgentMessages } = require('@librechat/agents');
|
||||
const { EModelEndpoint, ResourceType, PermissionBits } = require('librechat-data-provider');
|
||||
const {
|
||||
writeSSE,
|
||||
createRun,
|
||||
createChunk,
|
||||
buildToolSet,
|
||||
sendFinalChunk,
|
||||
createSafeUser,
|
||||
validateRequest,
|
||||
initializeAgent,
|
||||
getBalanceConfig,
|
||||
createErrorResponse,
|
||||
recordCollectedUsage,
|
||||
getTransactionsConfig,
|
||||
createToolExecuteHandler,
|
||||
buildNonStreamingResponse,
|
||||
createOpenAIStreamTracker,
|
||||
createOpenAIContentAggregator,
|
||||
isChatCompletionValidationFailure,
|
||||
} = require('@librechat/api');
|
||||
const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
|
||||
const { createToolEndCallback } = require('~/server/controllers/agents/callbacks');
|
||||
const { findAccessibleResources } = require('~/server/services/PermissionService');
|
||||
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
|
||||
const { getConvoFiles } = require('~/models/Conversation');
|
||||
const { getAgent, getAgents } = require('~/models/Agent');
|
||||
const db = require('~/models');
|
||||
|
||||
/**
|
||||
* Creates a tool loader function for the agent.
|
||||
* @param {AbortSignal} signal - The abort signal
|
||||
* @param {boolean} [definitionsOnly=true] - When true, returns only serializable
|
||||
* tool definitions without creating full tool instances (for event-driven mode)
|
||||
*/
|
||||
function createToolLoader(signal, definitionsOnly = true) {
|
||||
return async function loadTools({
|
||||
req,
|
||||
res,
|
||||
tools,
|
||||
model,
|
||||
agentId,
|
||||
provider,
|
||||
tool_options,
|
||||
tool_resources,
|
||||
}) {
|
||||
const agent = { id: agentId, tools, provider, model, tool_options };
|
||||
try {
|
||||
return await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
signal,
|
||||
tool_resources,
|
||||
definitionsOnly,
|
||||
streamId: null, // No resumable stream for OpenAI compat
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('Error loading tools for agent ' + agentId, error);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert content part to internal format
|
||||
* @param {Object} part - Content part
|
||||
* @returns {Object} Converted part
|
||||
*/
|
||||
function convertContentPart(part) {
|
||||
if (part.type === 'text') {
|
||||
return { type: 'text', text: part.text };
|
||||
}
|
||||
if (part.type === 'image_url') {
|
||||
return { type: 'image_url', image_url: part.image_url };
|
||||
}
|
||||
return part;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert OpenAI messages to internal format
|
||||
* @param {Array} messages - OpenAI format messages
|
||||
* @returns {Array} Internal format messages
|
||||
*/
|
||||
function convertMessages(messages) {
|
||||
return messages.map((msg) => {
|
||||
let content;
|
||||
if (typeof msg.content === 'string') {
|
||||
content = msg.content;
|
||||
} else if (msg.content) {
|
||||
content = msg.content.map(convertContentPart);
|
||||
} else {
|
||||
content = '';
|
||||
}
|
||||
|
||||
return {
|
||||
role: msg.role,
|
||||
content,
|
||||
...(msg.name && { name: msg.name }),
|
||||
...(msg.tool_calls && { tool_calls: msg.tool_calls }),
|
||||
...(msg.tool_call_id && { tool_call_id: msg.tool_call_id }),
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Send an error response in OpenAI format
|
||||
*/
|
||||
function sendErrorResponse(res, statusCode, message, type = 'invalid_request_error', code = null) {
|
||||
res.status(statusCode).json(createErrorResponse(message, type, code));
|
||||
}
|
||||
|
||||
/**
|
||||
* OpenAI-compatible chat completions controller for agents.
|
||||
*
|
||||
* POST /v1/chat/completions
|
||||
*
|
||||
* Request format:
|
||||
* {
|
||||
* "model": "agent_id_here",
|
||||
* "messages": [{"role": "user", "content": "Hello!"}],
|
||||
* "stream": true,
|
||||
* "conversation_id": "optional",
|
||||
* "parent_message_id": "optional"
|
||||
* }
|
||||
*/
|
||||
const OpenAIChatCompletionController = async (req, res) => {
|
||||
const appConfig = req.config;
|
||||
const requestStartTime = Date.now();
|
||||
|
||||
// Validate request
|
||||
const validation = validateRequest(req.body);
|
||||
if (isChatCompletionValidationFailure(validation)) {
|
||||
return sendErrorResponse(res, 400, validation.error);
|
||||
}
|
||||
|
||||
const request = validation.request;
|
||||
const agentId = request.model;
|
||||
|
||||
// Look up the agent
|
||||
const agent = await getAgent({ id: agentId });
|
||||
if (!agent) {
|
||||
return sendErrorResponse(
|
||||
res,
|
||||
404,
|
||||
`Agent not found: ${agentId}`,
|
||||
'invalid_request_error',
|
||||
'model_not_found',
|
||||
);
|
||||
}
|
||||
|
||||
// Generate IDs
|
||||
const requestId = `chatcmpl-${nanoid()}`;
|
||||
const conversationId = request.conversation_id ?? nanoid();
|
||||
const parentMessageId = request.parent_message_id ?? null;
|
||||
const created = Math.floor(Date.now() / 1000);
|
||||
|
||||
const context = {
|
||||
created,
|
||||
requestId,
|
||||
model: agentId,
|
||||
};
|
||||
|
||||
logger.debug(
|
||||
`[OpenAI API] Request ${requestId} started for agent ${agentId}, stream: ${request.stream}`,
|
||||
);
|
||||
|
||||
// Set up abort controller
|
||||
const abortController = new AbortController();
|
||||
|
||||
// Handle client disconnect
|
||||
req.on('close', () => {
|
||||
if (!abortController.signal.aborted) {
|
||||
abortController.abort();
|
||||
logger.debug('[OpenAI API] Client disconnected, aborting');
|
||||
}
|
||||
});
|
||||
|
||||
try {
|
||||
// Build allowed providers set
|
||||
const allowedProviders = new Set(
|
||||
appConfig?.endpoints?.[EModelEndpoint.agents]?.allowedProviders,
|
||||
);
|
||||
|
||||
// Create tool loader
|
||||
const loadTools = createToolLoader(abortController.signal);
|
||||
|
||||
// Initialize the agent first to check for disableStreaming
|
||||
const endpointOption = {
|
||||
endpoint: agent.provider,
|
||||
model_parameters: agent.model_parameters ?? {},
|
||||
};
|
||||
|
||||
const primaryConfig = await initializeAgent(
|
||||
{
|
||||
req,
|
||||
res,
|
||||
loadTools,
|
||||
requestFiles: [],
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
isInitialAgent: true,
|
||||
},
|
||||
{
|
||||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
// Determine if streaming is enabled (check both request and agent config)
|
||||
const streamingDisabled = !!primaryConfig.model_parameters?.disableStreaming;
|
||||
const isStreaming = request.stream === true && !streamingDisabled;
|
||||
|
||||
// Create tracker for streaming or aggregator for non-streaming
|
||||
const tracker = isStreaming ? createOpenAIStreamTracker() : null;
|
||||
const aggregator = isStreaming ? null : createOpenAIContentAggregator();
|
||||
|
||||
// Set up response for streaming
|
||||
if (isStreaming) {
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-cache');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
res.setHeader('X-Accel-Buffering', 'no');
|
||||
res.flushHeaders();
|
||||
|
||||
// Send initial chunk with role
|
||||
const initialChunk = createChunk(context, { role: 'assistant' });
|
||||
writeSSE(res, initialChunk);
|
||||
}
|
||||
|
||||
// Create handler config for OpenAI streaming (only used when streaming)
|
||||
const handlerConfig = isStreaming
|
||||
? {
|
||||
res,
|
||||
context,
|
||||
tracker,
|
||||
}
|
||||
: null;
|
||||
|
||||
const collectedUsage = [];
|
||||
/** @type {Promise<import('librechat-data-provider').TAttachment | null>[]} */
|
||||
const artifactPromises = [];
|
||||
|
||||
const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId: null });
|
||||
|
||||
const toolExecuteOptions = {
|
||||
loadTools: async (toolNames) => {
|
||||
return loadToolsForExecution({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
toolNames,
|
||||
signal: abortController.signal,
|
||||
toolRegistry: primaryConfig.toolRegistry,
|
||||
userMCPAuthMap: primaryConfig.userMCPAuthMap,
|
||||
tool_resources: primaryConfig.tool_resources,
|
||||
});
|
||||
},
|
||||
toolEndCallback,
|
||||
};
|
||||
|
||||
const openaiMessages = convertMessages(request.messages);
|
||||
|
||||
const toolSet = buildToolSet(primaryConfig);
|
||||
const { messages: formattedMessages, indexTokenCountMap } = formatAgentMessages(
|
||||
openaiMessages,
|
||||
{},
|
||||
toolSet,
|
||||
);
|
||||
|
||||
/**
|
||||
* Create a simple handler that processes data
|
||||
*/
|
||||
const createHandler = (processor) => ({
|
||||
handle: (_event, data) => {
|
||||
if (processor) {
|
||||
processor(data);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
/**
|
||||
* Stream text content in OpenAI format
|
||||
*/
|
||||
const streamText = (text) => {
|
||||
if (!text) {
|
||||
return;
|
||||
}
|
||||
if (isStreaming) {
|
||||
tracker.addText();
|
||||
writeSSE(res, createChunk(context, { content: text }));
|
||||
} else {
|
||||
aggregator.addText(text);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Stream reasoning content in OpenAI format (OpenRouter convention)
|
||||
*/
|
||||
const streamReasoning = (text) => {
|
||||
if (!text) {
|
||||
return;
|
||||
}
|
||||
if (isStreaming) {
|
||||
tracker.addReasoning();
|
||||
writeSSE(res, createChunk(context, { reasoning: text }));
|
||||
} else {
|
||||
aggregator.addReasoning(text);
|
||||
}
|
||||
};
|
||||
|
||||
// Event handlers for OpenAI-compatible streaming
|
||||
const handlers = {
|
||||
// Text content streaming
|
||||
on_message_delta: createHandler((data) => {
|
||||
const content = data?.delta?.content;
|
||||
if (Array.isArray(content)) {
|
||||
for (const part of content) {
|
||||
if (part.type === 'text' && part.text) {
|
||||
streamText(part.text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}),
|
||||
|
||||
// Reasoning/thinking content streaming
|
||||
on_reasoning_delta: createHandler((data) => {
|
||||
const content = data?.delta?.content;
|
||||
if (Array.isArray(content)) {
|
||||
for (const part of content) {
|
||||
const text = part.think || part.text;
|
||||
if (text) {
|
||||
streamReasoning(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}),
|
||||
|
||||
// Tool call initiation - streams id and name (from on_run_step)
|
||||
on_run_step: createHandler((data) => {
|
||||
const stepDetails = data?.stepDetails;
|
||||
if (stepDetails?.type === 'tool_calls' && stepDetails.tool_calls) {
|
||||
for (const tc of stepDetails.tool_calls) {
|
||||
const toolIndex = data.index ?? 0;
|
||||
const toolId = tc.id ?? '';
|
||||
const toolName = tc.name ?? '';
|
||||
const toolCall = {
|
||||
id: toolId,
|
||||
type: 'function',
|
||||
function: { name: toolName, arguments: '' },
|
||||
};
|
||||
|
||||
// Track tool call in tracker or aggregator
|
||||
if (isStreaming) {
|
||||
if (!tracker.toolCalls.has(toolIndex)) {
|
||||
tracker.toolCalls.set(toolIndex, toolCall);
|
||||
}
|
||||
// Stream initial tool call chunk (like OpenAI does)
|
||||
writeSSE(
|
||||
res,
|
||||
createChunk(context, {
|
||||
tool_calls: [{ index: toolIndex, ...toolCall }],
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
if (!aggregator.toolCalls.has(toolIndex)) {
|
||||
aggregator.toolCalls.set(toolIndex, toolCall);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}),
|
||||
|
||||
// Tool call argument streaming (from on_run_step_delta)
|
||||
on_run_step_delta: createHandler((data) => {
|
||||
const delta = data?.delta;
|
||||
if (delta?.type === 'tool_calls' && delta.tool_calls) {
|
||||
for (const tc of delta.tool_calls) {
|
||||
const args = tc.args ?? '';
|
||||
if (!args) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const toolIndex = tc.index ?? 0;
|
||||
|
||||
// Update tool call arguments
|
||||
const targetMap = isStreaming ? tracker.toolCalls : aggregator.toolCalls;
|
||||
const tracked = targetMap.get(toolIndex);
|
||||
if (tracked) {
|
||||
tracked.function.arguments += args;
|
||||
}
|
||||
|
||||
// Stream argument delta (only for streaming)
|
||||
if (isStreaming) {
|
||||
writeSSE(
|
||||
res,
|
||||
createChunk(context, {
|
||||
tool_calls: [
|
||||
{
|
||||
index: toolIndex,
|
||||
function: { arguments: args },
|
||||
},
|
||||
],
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}),
|
||||
|
||||
// Usage tracking
|
||||
on_chat_model_end: createHandler((data) => {
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
collectedUsage.push(usage);
|
||||
const target = isStreaming ? tracker : aggregator;
|
||||
target.usage.promptTokens += usage.input_tokens ?? 0;
|
||||
target.usage.completionTokens += usage.output_tokens ?? 0;
|
||||
}
|
||||
}),
|
||||
on_run_step_completed: createHandler(),
|
||||
// Use proper ToolEndHandler for processing artifacts (images, file citations, code output)
|
||||
on_tool_end: new ToolEndHandler(toolEndCallback, logger),
|
||||
on_chain_stream: createHandler(),
|
||||
on_chain_end: createHandler(),
|
||||
on_agent_update: createHandler(),
|
||||
on_custom_event: createHandler(),
|
||||
// Event-driven tool execution handler
|
||||
on_tool_execute: createToolExecuteHandler(toolExecuteOptions),
|
||||
};
|
||||
|
||||
// Create and run the agent
|
||||
const userId = req.user?.id ?? 'api-user';
|
||||
|
||||
// Extract userMCPAuthMap from primaryConfig (needed for MCP tool connections)
|
||||
const userMCPAuthMap = primaryConfig.userMCPAuthMap;
|
||||
|
||||
const run = await createRun({
|
||||
agents: [primaryConfig],
|
||||
messages: formattedMessages,
|
||||
indexTokenCountMap,
|
||||
runId: requestId,
|
||||
signal: abortController.signal,
|
||||
customHandlers: handlers,
|
||||
requestBody: {
|
||||
messageId: requestId,
|
||||
conversationId,
|
||||
},
|
||||
user: { id: userId },
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
throw new Error('Failed to create agent run');
|
||||
}
|
||||
|
||||
// Process the stream
|
||||
const config = {
|
||||
runName: 'AgentRun',
|
||||
configurable: {
|
||||
thread_id: conversationId,
|
||||
user_id: userId,
|
||||
user: createSafeUser(req.user),
|
||||
...(userMCPAuthMap != null && { userMCPAuthMap }),
|
||||
},
|
||||
signal: abortController.signal,
|
||||
streamMode: 'values',
|
||||
version: 'v2',
|
||||
};
|
||||
|
||||
await run.processStream({ messages: formattedMessages }, config, {
|
||||
callbacks: {
|
||||
[Callback.TOOL_ERROR]: (graph, error, toolId) => {
|
||||
logger.error(`[OpenAI API] Tool Error "${toolId}"`, error);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Record token usage against balance
|
||||
const balanceConfig = getBalanceConfig(appConfig);
|
||||
const transactionsConfig = getTransactionsConfig(appConfig);
|
||||
recordCollectedUsage(
|
||||
{ spendTokens, spendStructuredTokens },
|
||||
{
|
||||
user: userId,
|
||||
conversationId,
|
||||
collectedUsage,
|
||||
context: 'message',
|
||||
balance: balanceConfig,
|
||||
transactions: transactionsConfig,
|
||||
model: primaryConfig.model || agent.model_parameters?.model,
|
||||
},
|
||||
).catch((err) => {
|
||||
logger.error('[OpenAI API] Error recording usage:', err);
|
||||
});
|
||||
|
||||
// Finalize response
|
||||
const duration = Date.now() - requestStartTime;
|
||||
if (isStreaming) {
|
||||
sendFinalChunk(handlerConfig);
|
||||
res.end();
|
||||
logger.debug(`[OpenAI API] Request ${requestId} completed in ${duration}ms (streaming)`);
|
||||
|
||||
// Wait for artifact processing after response ends (non-blocking)
|
||||
if (artifactPromises.length > 0) {
|
||||
Promise.all(artifactPromises).catch((artifactError) => {
|
||||
logger.warn('[OpenAI API] Error processing artifacts:', artifactError);
|
||||
});
|
||||
}
|
||||
} else {
|
||||
// For non-streaming, wait for artifacts before sending response
|
||||
if (artifactPromises.length > 0) {
|
||||
try {
|
||||
await Promise.all(artifactPromises);
|
||||
} catch (artifactError) {
|
||||
logger.warn('[OpenAI API] Error processing artifacts:', artifactError);
|
||||
}
|
||||
}
|
||||
|
||||
// Build usage from aggregated data
|
||||
const usage = {
|
||||
prompt_tokens: aggregator.usage.promptTokens,
|
||||
completion_tokens: aggregator.usage.completionTokens,
|
||||
total_tokens: aggregator.usage.promptTokens + aggregator.usage.completionTokens,
|
||||
};
|
||||
|
||||
if (aggregator.usage.reasoningTokens > 0) {
|
||||
usage.completion_tokens_details = {
|
||||
reasoning_tokens: aggregator.usage.reasoningTokens,
|
||||
};
|
||||
}
|
||||
|
||||
const response = buildNonStreamingResponse(
|
||||
context,
|
||||
aggregator.getText(),
|
||||
aggregator.getReasoning(),
|
||||
aggregator.toolCalls,
|
||||
usage,
|
||||
);
|
||||
res.json(response);
|
||||
logger.debug(`[OpenAI API] Request ${requestId} completed in ${duration}ms (non-streaming)`);
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'An error occurred';
|
||||
logger.error('[OpenAI API] Error:', error);
|
||||
|
||||
// Check if we already started streaming (headers sent)
|
||||
if (res.headersSent) {
|
||||
// Headers already sent, send error in stream
|
||||
const errorChunk = createChunk(context, { content: `\n\nError: ${errorMessage}` }, 'stop');
|
||||
writeSSE(res, errorChunk);
|
||||
writeSSE(res, '[DONE]');
|
||||
res.end();
|
||||
} else {
|
||||
// Forward upstream provider status codes (e.g., Anthropic 400s) instead of masking as 500
|
||||
const statusCode =
|
||||
typeof error?.status === 'number' && error.status >= 400 && error.status < 600
|
||||
? error.status
|
||||
: 500;
|
||||
const errorType =
|
||||
statusCode >= 400 && statusCode < 500 ? 'invalid_request_error' : 'server_error';
|
||||
sendErrorResponse(res, statusCode, errorMessage, errorType);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* List available agents as models (filtered by remote access permissions)
|
||||
*
|
||||
* GET /v1/models
|
||||
*/
|
||||
const ListModelsController = async (req, res) => {
|
||||
try {
|
||||
const userId = req.user?.id;
|
||||
const userRole = req.user?.role;
|
||||
|
||||
if (!userId) {
|
||||
return sendErrorResponse(res, 401, 'Authentication required', 'auth_error');
|
||||
}
|
||||
|
||||
// Find agents the user has remote access to (VIEW permission on REMOTE_AGENT)
|
||||
const accessibleAgentIds = await findAccessibleResources({
|
||||
userId,
|
||||
role: userRole,
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
requiredPermissions: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
// Get the accessible agents
|
||||
let agents = [];
|
||||
if (accessibleAgentIds.length > 0) {
|
||||
agents = await getAgents({ _id: { $in: accessibleAgentIds } });
|
||||
}
|
||||
|
||||
const models = agents.map((agent) => ({
|
||||
id: agent.id,
|
||||
object: 'model',
|
||||
created: Math.floor(new Date(agent.createdAt || Date.now()).getTime() / 1000),
|
||||
owned_by: 'librechat',
|
||||
permission: [],
|
||||
root: agent.id,
|
||||
parent: null,
|
||||
// LibreChat extensions
|
||||
name: agent.name,
|
||||
description: agent.description,
|
||||
provider: agent.provider,
|
||||
}));
|
||||
|
||||
res.json({
|
||||
object: 'list',
|
||||
data: models,
|
||||
});
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Failed to list models';
|
||||
logger.error('[OpenAI API] Error listing models:', error);
|
||||
sendErrorResponse(res, 500, errorMessage, 'server_error');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Get a specific model/agent (with remote access permission check)
|
||||
*
|
||||
* GET /v1/models/:model
|
||||
*/
|
||||
const GetModelController = async (req, res) => {
|
||||
try {
|
||||
const { model } = req.params;
|
||||
const userId = req.user?.id;
|
||||
const userRole = req.user?.role;
|
||||
|
||||
if (!userId) {
|
||||
return sendErrorResponse(res, 401, 'Authentication required', 'auth_error');
|
||||
}
|
||||
|
||||
const agent = await getAgent({ id: model });
|
||||
|
||||
if (!agent) {
|
||||
return sendErrorResponse(
|
||||
res,
|
||||
404,
|
||||
`Model not found: ${model}`,
|
||||
'invalid_request_error',
|
||||
'model_not_found',
|
||||
);
|
||||
}
|
||||
|
||||
// Check if user has remote access to this agent
|
||||
const accessibleAgentIds = await findAccessibleResources({
|
||||
userId,
|
||||
role: userRole,
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
requiredPermissions: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
const hasAccess = accessibleAgentIds.some((id) => id.toString() === agent._id.toString());
|
||||
|
||||
if (!hasAccess) {
|
||||
return sendErrorResponse(
|
||||
res,
|
||||
403,
|
||||
`No remote access to model: ${model}`,
|
||||
'permission_error',
|
||||
'access_denied',
|
||||
);
|
||||
}
|
||||
|
||||
res.json({
|
||||
id: agent.id,
|
||||
object: 'model',
|
||||
created: Math.floor(new Date(agent.createdAt || Date.now()).getTime() / 1000),
|
||||
owned_by: 'librechat',
|
||||
permission: [],
|
||||
root: agent.id,
|
||||
parent: null,
|
||||
// LibreChat extensions
|
||||
name: agent.name,
|
||||
description: agent.description,
|
||||
provider: agent.provider,
|
||||
});
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Failed to get model';
|
||||
logger.error('[OpenAI API] Error getting model:', error);
|
||||
sendErrorResponse(res, 500, errorMessage, 'server_error');
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
OpenAIChatCompletionController,
|
||||
ListModelsController,
|
||||
GetModelController,
|
||||
};
|
||||
|
|
@ -67,7 +67,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
let client = null;
|
||||
|
||||
try {
|
||||
logger.debug(`[ResumableAgentController] Creating job`, {
|
||||
streamId,
|
||||
conversationId,
|
||||
reqConversationId,
|
||||
userId,
|
||||
});
|
||||
|
||||
const job = await GenerationJobManager.createJob(streamId, userId, conversationId);
|
||||
const jobCreatedAt = job.createdAt; // Capture creation time to detect job replacement
|
||||
req._resumableStreamId = streamId;
|
||||
|
||||
// Send JSON response IMMEDIATELY so client can connect to SSE stream
|
||||
|
|
@ -272,6 +280,33 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
});
|
||||
}
|
||||
|
||||
// CRITICAL: Save response message BEFORE emitting final event.
|
||||
// This prevents race conditions where the client sends a follow-up message
|
||||
// before the response is saved to the database, causing orphaned parentMessageIds.
|
||||
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
|
||||
await saveMessage(
|
||||
req,
|
||||
{ ...response, user: userId, unfinished: wasAbortedBeforeComplete },
|
||||
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
|
||||
);
|
||||
}
|
||||
|
||||
// Check if our job was replaced by a new request before emitting
|
||||
// This prevents stale requests from emitting events to newer jobs
|
||||
const currentJob = await GenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== jobCreatedAt;
|
||||
|
||||
if (jobWasReplaced) {
|
||||
logger.debug(`[ResumableAgentController] Skipping FINAL emit - job was replaced`, {
|
||||
streamId,
|
||||
originalCreatedAt: jobCreatedAt,
|
||||
currentCreatedAt: currentJob?.createdAt,
|
||||
});
|
||||
// Still decrement pending request since we incremented at start
|
||||
await decrementPendingRequest(userId);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!wasAbortedBeforeComplete) {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
|
|
@ -281,27 +316,35 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
responseMessage: { ...response },
|
||||
};
|
||||
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
logger.debug(`[ResumableAgentController] Emitting FINAL event`, {
|
||||
streamId,
|
||||
wasAbortedBeforeComplete,
|
||||
userMessageId: userMessage?.messageId,
|
||||
responseMessageId: response?.messageId,
|
||||
conversationId: conversation?.conversationId,
|
||||
});
|
||||
|
||||
await GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId);
|
||||
await decrementPendingRequest(userId);
|
||||
|
||||
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
|
||||
await saveMessage(
|
||||
req,
|
||||
{ ...response, user: userId },
|
||||
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
|
||||
);
|
||||
}
|
||||
} else {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
conversation,
|
||||
title: conversation.title,
|
||||
requestMessage: sanitizeMessageForTransmit(userMessage),
|
||||
responseMessage: { ...response, error: true },
|
||||
error: { message: 'Request was aborted' },
|
||||
responseMessage: { ...response, unfinished: true },
|
||||
};
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
|
||||
logger.debug(`[ResumableAgentController] Emitting ABORTED FINAL event`, {
|
||||
streamId,
|
||||
wasAbortedBeforeComplete,
|
||||
userMessageId: userMessage?.messageId,
|
||||
responseMessageId: response?.messageId,
|
||||
conversationId: conversation?.conversationId,
|
||||
});
|
||||
|
||||
await GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId, 'Request aborted');
|
||||
await decrementPendingRequest(userId);
|
||||
}
|
||||
|
|
@ -334,7 +377,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
// abortJob already handled emitDone and completeJob
|
||||
} else {
|
||||
logger.error(`[ResumableAgentController] Generation error for ${streamId}:`, error);
|
||||
GenerationJobManager.emitError(streamId, error.message || 'Generation failed');
|
||||
await GenerationJobManager.emitError(streamId, error.message || 'Generation failed');
|
||||
GenerationJobManager.completeJob(streamId, error.message);
|
||||
}
|
||||
|
||||
|
|
@ -363,7 +406,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
res.status(500).json({ error: error.message || 'Failed to start generation' });
|
||||
} else {
|
||||
// JSON already sent, emit error to stream so client can receive it
|
||||
GenerationJobManager.emitError(streamId, error.message || 'Failed to start generation');
|
||||
await GenerationJobManager.emitError(streamId, error.message || 'Failed to start generation');
|
||||
}
|
||||
GenerationJobManager.completeJob(streamId, error.message);
|
||||
await decrementPendingRequest(userId);
|
||||
|
|
|
|||
889
api/server/controllers/agents/responses.js
Normal file
889
api/server/controllers/agents/responses.js
Normal file
|
|
@ -0,0 +1,889 @@
|
|||
const { nanoid } = require('nanoid');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Callback, ToolEndHandler, formatAgentMessages } = require('@librechat/agents');
|
||||
const { EModelEndpoint, ResourceType, PermissionBits } = require('librechat-data-provider');
|
||||
const {
|
||||
createRun,
|
||||
buildToolSet,
|
||||
createSafeUser,
|
||||
initializeAgent,
|
||||
getBalanceConfig,
|
||||
recordCollectedUsage,
|
||||
getTransactionsConfig,
|
||||
createToolExecuteHandler,
|
||||
// Responses API
|
||||
writeDone,
|
||||
buildResponse,
|
||||
generateResponseId,
|
||||
isValidationFailure,
|
||||
emitResponseCreated,
|
||||
createResponseContext,
|
||||
createResponseTracker,
|
||||
setupStreamingResponse,
|
||||
emitResponseInProgress,
|
||||
convertInputToMessages,
|
||||
validateResponseRequest,
|
||||
buildAggregatedResponse,
|
||||
createResponseAggregator,
|
||||
sendResponsesErrorResponse,
|
||||
createResponsesEventHandlers,
|
||||
createAggregatorEventHandlers,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
createResponsesToolEndCallback,
|
||||
createToolEndCallback,
|
||||
} = require('~/server/controllers/agents/callbacks');
|
||||
const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
|
||||
const { findAccessibleResources } = require('~/server/services/PermissionService');
|
||||
const { getConvoFiles, saveConvo, getConvo } = require('~/models/Conversation');
|
||||
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
|
||||
const { getAgent, getAgents } = require('~/models/Agent');
|
||||
const db = require('~/models');
|
||||
|
||||
/** @type {import('@librechat/api').AppConfig | null} */
|
||||
let appConfig = null;
|
||||
|
||||
/**
|
||||
* Set the app config for the controller
|
||||
* @param {import('@librechat/api').AppConfig} config
|
||||
*/
|
||||
function setAppConfig(config) {
|
||||
appConfig = config;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a tool loader function for the agent.
|
||||
* @param {AbortSignal} signal - The abort signal
|
||||
* @param {boolean} [definitionsOnly=true] - When true, returns only serializable
|
||||
* tool definitions without creating full tool instances (for event-driven mode)
|
||||
*/
|
||||
function createToolLoader(signal, definitionsOnly = true) {
|
||||
return async function loadTools({
|
||||
req,
|
||||
res,
|
||||
tools,
|
||||
model,
|
||||
agentId,
|
||||
provider,
|
||||
tool_options,
|
||||
tool_resources,
|
||||
}) {
|
||||
const agent = { id: agentId, tools, provider, model, tool_options };
|
||||
try {
|
||||
return await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
signal,
|
||||
tool_resources,
|
||||
definitionsOnly,
|
||||
streamId: null,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('Error loading tools for agent ' + agentId, error);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert Open Responses input items to internal messages
|
||||
* @param {import('@librechat/api').InputItem[]} input
|
||||
* @returns {Array} Internal messages
|
||||
*/
|
||||
function convertToInternalMessages(input) {
|
||||
return convertInputToMessages(input);
|
||||
}
|
||||
|
||||
/**
|
||||
* Load messages from a previous response/conversation
|
||||
* @param {string} conversationId - The conversation/response ID
|
||||
* @param {string} userId - The user ID
|
||||
* @returns {Promise<Array>} Messages from the conversation
|
||||
*/
|
||||
async function loadPreviousMessages(conversationId, userId) {
|
||||
try {
|
||||
const messages = await db.getMessages({ conversationId, user: userId });
|
||||
if (!messages || messages.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
// Convert stored messages to internal format
|
||||
return messages.map((msg) => {
|
||||
const internalMsg = {
|
||||
role: msg.isCreatedByUser ? 'user' : 'assistant',
|
||||
content: '',
|
||||
messageId: msg.messageId,
|
||||
};
|
||||
|
||||
// Handle content - could be string or array
|
||||
if (typeof msg.text === 'string') {
|
||||
internalMsg.content = msg.text;
|
||||
} else if (Array.isArray(msg.content)) {
|
||||
// Handle content parts
|
||||
internalMsg.content = msg.content;
|
||||
} else if (msg.text) {
|
||||
internalMsg.content = String(msg.text);
|
||||
}
|
||||
|
||||
return internalMsg;
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[Responses API] Error loading previous messages:', error);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Save input messages to database
|
||||
* @param {import('express').Request} req
|
||||
* @param {string} conversationId
|
||||
* @param {Array} inputMessages - Internal format messages
|
||||
* @param {string} agentId
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function saveInputMessages(req, conversationId, inputMessages, agentId) {
|
||||
for (const msg of inputMessages) {
|
||||
if (msg.role === 'user') {
|
||||
await db.saveMessage(
|
||||
req,
|
||||
{
|
||||
messageId: msg.messageId || nanoid(),
|
||||
conversationId,
|
||||
parentMessageId: null,
|
||||
isCreatedByUser: true,
|
||||
text: typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content),
|
||||
sender: 'User',
|
||||
endpoint: EModelEndpoint.agents,
|
||||
model: agentId,
|
||||
},
|
||||
{ context: 'Responses API - save user input' },
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Save response output to database
|
||||
* @param {import('express').Request} req
|
||||
* @param {string} conversationId
|
||||
* @param {string} responseId
|
||||
* @param {import('@librechat/api').Response} response
|
||||
* @param {string} agentId
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function saveResponseOutput(req, conversationId, responseId, response, agentId) {
|
||||
// Extract text content from output items
|
||||
let responseText = '';
|
||||
for (const item of response.output) {
|
||||
if (item.type === 'message' && item.content) {
|
||||
for (const part of item.content) {
|
||||
if (part.type === 'output_text' && part.text) {
|
||||
responseText += part.text;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save the assistant message
|
||||
await db.saveMessage(
|
||||
req,
|
||||
{
|
||||
messageId: responseId,
|
||||
conversationId,
|
||||
parentMessageId: null,
|
||||
isCreatedByUser: false,
|
||||
text: responseText,
|
||||
sender: 'Agent',
|
||||
endpoint: EModelEndpoint.agents,
|
||||
model: agentId,
|
||||
finish_reason: response.status === 'completed' ? 'stop' : response.status,
|
||||
tokenCount: response.usage?.output_tokens,
|
||||
},
|
||||
{ context: 'Responses API - save assistant response' },
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Save or update conversation
|
||||
* @param {import('express').Request} req
|
||||
* @param {string} conversationId
|
||||
* @param {string} agentId
|
||||
* @param {object} agent
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function saveConversation(req, conversationId, agentId, agent) {
|
||||
await saveConvo(
|
||||
req,
|
||||
{
|
||||
conversationId,
|
||||
endpoint: EModelEndpoint.agents,
|
||||
agentId,
|
||||
title: agent?.name || 'Open Responses Conversation',
|
||||
model: agent?.model,
|
||||
},
|
||||
{ context: 'Responses API - save conversation' },
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert stored messages to Open Responses output format
|
||||
* @param {Array} messages - Stored messages
|
||||
* @returns {Array} Output items
|
||||
*/
|
||||
function convertMessagesToOutputItems(messages) {
|
||||
const output = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
if (!msg.isCreatedByUser) {
|
||||
output.push({
|
||||
type: 'message',
|
||||
id: msg.messageId,
|
||||
role: 'assistant',
|
||||
status: 'completed',
|
||||
content: [
|
||||
{
|
||||
type: 'output_text',
|
||||
text: msg.text || '',
|
||||
annotations: [],
|
||||
},
|
||||
],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create Response - POST /v1/responses
|
||||
*
|
||||
* Creates a model response following the Open Responses API specification.
|
||||
* Supports both streaming and non-streaming responses.
|
||||
*
|
||||
* @param {import('express').Request} req
|
||||
* @param {import('express').Response} res
|
||||
*/
|
||||
const createResponse = async (req, res) => {
|
||||
const requestStartTime = Date.now();
|
||||
|
||||
// Validate request
|
||||
const validation = validateResponseRequest(req.body);
|
||||
if (isValidationFailure(validation)) {
|
||||
return sendResponsesErrorResponse(res, 400, validation.error);
|
||||
}
|
||||
|
||||
const request = validation.request;
|
||||
const agentId = request.model;
|
||||
const isStreaming = request.stream === true;
|
||||
|
||||
// Look up the agent
|
||||
const agent = await getAgent({ id: agentId });
|
||||
if (!agent) {
|
||||
return sendResponsesErrorResponse(
|
||||
res,
|
||||
404,
|
||||
`Agent not found: ${agentId}`,
|
||||
'not_found',
|
||||
'model_not_found',
|
||||
);
|
||||
}
|
||||
|
||||
// Generate IDs
|
||||
const responseId = generateResponseId();
|
||||
const conversationId = request.previous_response_id ?? uuidv4();
|
||||
const parentMessageId = null;
|
||||
|
||||
// Create response context
|
||||
const context = createResponseContext(request, responseId);
|
||||
|
||||
logger.debug(
|
||||
`[Responses API] Request ${responseId} started for agent ${agentId}, stream: ${isStreaming}`,
|
||||
);
|
||||
|
||||
// Set up abort controller
|
||||
const abortController = new AbortController();
|
||||
|
||||
// Handle client disconnect
|
||||
req.on('close', () => {
|
||||
if (!abortController.signal.aborted) {
|
||||
abortController.abort();
|
||||
logger.debug('[Responses API] Client disconnected, aborting');
|
||||
}
|
||||
});
|
||||
|
||||
try {
|
||||
// Build allowed providers set
|
||||
const allowedProviders = new Set(
|
||||
appConfig?.endpoints?.[EModelEndpoint.agents]?.allowedProviders,
|
||||
);
|
||||
|
||||
// Create tool loader
|
||||
const loadTools = createToolLoader(abortController.signal);
|
||||
|
||||
// Initialize the agent first to check for disableStreaming
|
||||
const endpointOption = {
|
||||
endpoint: agent.provider,
|
||||
model_parameters: agent.model_parameters ?? {},
|
||||
};
|
||||
|
||||
const primaryConfig = await initializeAgent(
|
||||
{
|
||||
req,
|
||||
res,
|
||||
loadTools,
|
||||
requestFiles: [],
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
isInitialAgent: true,
|
||||
},
|
||||
{
|
||||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
// Determine if streaming is enabled (check both request and agent config)
|
||||
const streamingDisabled = !!primaryConfig.model_parameters?.disableStreaming;
|
||||
const actuallyStreaming = isStreaming && !streamingDisabled;
|
||||
|
||||
// Load previous messages if previous_response_id is provided
|
||||
let previousMessages = [];
|
||||
if (request.previous_response_id) {
|
||||
const userId = req.user?.id ?? 'api-user';
|
||||
previousMessages = await loadPreviousMessages(request.previous_response_id, userId);
|
||||
}
|
||||
|
||||
// Convert input to internal messages
|
||||
const inputMessages = convertToInternalMessages(
|
||||
typeof request.input === 'string' ? request.input : request.input,
|
||||
);
|
||||
|
||||
// Merge previous messages with new input
|
||||
const allMessages = [...previousMessages, ...inputMessages];
|
||||
|
||||
const toolSet = buildToolSet(primaryConfig);
|
||||
const { messages: formattedMessages, indexTokenCountMap } = formatAgentMessages(
|
||||
allMessages,
|
||||
{},
|
||||
toolSet,
|
||||
);
|
||||
|
||||
// Create tracker for streaming or aggregator for non-streaming
|
||||
const tracker = actuallyStreaming ? createResponseTracker() : null;
|
||||
const aggregator = actuallyStreaming ? null : createResponseAggregator();
|
||||
|
||||
// Set up response for streaming
|
||||
if (actuallyStreaming) {
|
||||
setupStreamingResponse(res);
|
||||
|
||||
// Create handler config
|
||||
const handlerConfig = {
|
||||
res,
|
||||
context,
|
||||
tracker,
|
||||
};
|
||||
|
||||
// Emit response.created then response.in_progress per Open Responses spec
|
||||
emitResponseCreated(handlerConfig);
|
||||
emitResponseInProgress(handlerConfig);
|
||||
|
||||
// Create event handlers
|
||||
const { handlers: responsesHandlers, finalizeStream } =
|
||||
createResponsesEventHandlers(handlerConfig);
|
||||
|
||||
// Collect usage for balance tracking
|
||||
const collectedUsage = [];
|
||||
|
||||
// Artifact promises for processing tool outputs
|
||||
/** @type {Promise<import('librechat-data-provider').TAttachment | null>[]} */
|
||||
const artifactPromises = [];
|
||||
// Use Responses API-specific callback that emits librechat:attachment events
|
||||
const toolEndCallback = createResponsesToolEndCallback({
|
||||
req,
|
||||
res,
|
||||
tracker,
|
||||
artifactPromises,
|
||||
});
|
||||
|
||||
// Create tool execute options for event-driven tool execution
|
||||
const toolExecuteOptions = {
|
||||
loadTools: async (toolNames) => {
|
||||
return loadToolsForExecution({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
toolNames,
|
||||
signal: abortController.signal,
|
||||
toolRegistry: primaryConfig.toolRegistry,
|
||||
userMCPAuthMap: primaryConfig.userMCPAuthMap,
|
||||
tool_resources: primaryConfig.tool_resources,
|
||||
});
|
||||
},
|
||||
toolEndCallback,
|
||||
};
|
||||
|
||||
// Combine handlers
|
||||
const handlers = {
|
||||
on_message_delta: responsesHandlers.on_message_delta,
|
||||
on_reasoning_delta: responsesHandlers.on_reasoning_delta,
|
||||
on_run_step: responsesHandlers.on_run_step,
|
||||
on_run_step_delta: responsesHandlers.on_run_step_delta,
|
||||
on_chat_model_end: {
|
||||
handle: (event, data) => {
|
||||
responsesHandlers.on_chat_model_end.handle(event, data);
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
collectedUsage.push(usage);
|
||||
}
|
||||
},
|
||||
},
|
||||
on_tool_end: new ToolEndHandler(toolEndCallback, logger),
|
||||
on_run_step_completed: { handle: () => {} },
|
||||
on_chain_stream: { handle: () => {} },
|
||||
on_chain_end: { handle: () => {} },
|
||||
on_agent_update: { handle: () => {} },
|
||||
on_custom_event: { handle: () => {} },
|
||||
on_tool_execute: createToolExecuteHandler(toolExecuteOptions),
|
||||
};
|
||||
|
||||
// Create and run the agent
|
||||
const userId = req.user?.id ?? 'api-user';
|
||||
const userMCPAuthMap = primaryConfig.userMCPAuthMap;
|
||||
|
||||
const run = await createRun({
|
||||
agents: [primaryConfig],
|
||||
messages: formattedMessages,
|
||||
indexTokenCountMap,
|
||||
runId: responseId,
|
||||
signal: abortController.signal,
|
||||
customHandlers: handlers,
|
||||
requestBody: {
|
||||
messageId: responseId,
|
||||
conversationId,
|
||||
},
|
||||
user: { id: userId },
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
throw new Error('Failed to create agent run');
|
||||
}
|
||||
|
||||
// Process the stream
|
||||
const config = {
|
||||
runName: 'AgentRun',
|
||||
configurable: {
|
||||
thread_id: conversationId,
|
||||
user_id: userId,
|
||||
user: createSafeUser(req.user),
|
||||
...(userMCPAuthMap != null && { userMCPAuthMap }),
|
||||
},
|
||||
signal: abortController.signal,
|
||||
streamMode: 'values',
|
||||
version: 'v2',
|
||||
};
|
||||
|
||||
await run.processStream({ messages: formattedMessages }, config, {
|
||||
callbacks: {
|
||||
[Callback.TOOL_ERROR]: (graph, error, toolId) => {
|
||||
logger.error(`[Responses API] Tool Error "${toolId}"`, error);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Record token usage against balance
|
||||
const balanceConfig = getBalanceConfig(req.config);
|
||||
const transactionsConfig = getTransactionsConfig(req.config);
|
||||
recordCollectedUsage(
|
||||
{ spendTokens, spendStructuredTokens },
|
||||
{
|
||||
user: userId,
|
||||
conversationId,
|
||||
collectedUsage,
|
||||
context: 'message',
|
||||
balance: balanceConfig,
|
||||
transactions: transactionsConfig,
|
||||
model: primaryConfig.model || agent.model_parameters?.model,
|
||||
},
|
||||
).catch((err) => {
|
||||
logger.error('[Responses API] Error recording usage:', err);
|
||||
});
|
||||
|
||||
// Finalize the stream
|
||||
finalizeStream();
|
||||
res.end();
|
||||
|
||||
const duration = Date.now() - requestStartTime;
|
||||
logger.debug(`[Responses API] Request ${responseId} completed in ${duration}ms (streaming)`);
|
||||
|
||||
// Save to database if store: true
|
||||
if (request.store === true) {
|
||||
try {
|
||||
// Save conversation
|
||||
await saveConversation(req, conversationId, agentId, agent);
|
||||
|
||||
// Save input messages
|
||||
await saveInputMessages(req, conversationId, inputMessages, agentId);
|
||||
|
||||
// Build response for saving (use tracker with buildResponse for streaming)
|
||||
const finalResponse = buildResponse(context, tracker, 'completed');
|
||||
await saveResponseOutput(req, conversationId, responseId, finalResponse, agentId);
|
||||
|
||||
logger.debug(
|
||||
`[Responses API] Stored response ${responseId} in conversation ${conversationId}`,
|
||||
);
|
||||
} catch (saveError) {
|
||||
logger.error('[Responses API] Error saving response:', saveError);
|
||||
// Don't fail the request if saving fails
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for artifact processing after response ends (non-blocking)
|
||||
if (artifactPromises.length > 0) {
|
||||
Promise.all(artifactPromises).catch((artifactError) => {
|
||||
logger.warn('[Responses API] Error processing artifacts:', artifactError);
|
||||
});
|
||||
}
|
||||
} else {
|
||||
const aggregatorHandlers = createAggregatorEventHandlers(aggregator);
|
||||
|
||||
// Collect usage for balance tracking
|
||||
const collectedUsage = [];
|
||||
|
||||
/** @type {Promise<import('librechat-data-provider').TAttachment | null>[]} */
|
||||
const artifactPromises = [];
|
||||
const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId: null });
|
||||
|
||||
const toolExecuteOptions = {
|
||||
loadTools: async (toolNames) => {
|
||||
return loadToolsForExecution({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
toolNames,
|
||||
signal: abortController.signal,
|
||||
toolRegistry: primaryConfig.toolRegistry,
|
||||
userMCPAuthMap: primaryConfig.userMCPAuthMap,
|
||||
tool_resources: primaryConfig.tool_resources,
|
||||
});
|
||||
},
|
||||
toolEndCallback,
|
||||
};
|
||||
|
||||
const handlers = {
|
||||
on_message_delta: aggregatorHandlers.on_message_delta,
|
||||
on_reasoning_delta: aggregatorHandlers.on_reasoning_delta,
|
||||
on_run_step: aggregatorHandlers.on_run_step,
|
||||
on_run_step_delta: aggregatorHandlers.on_run_step_delta,
|
||||
on_chat_model_end: {
|
||||
handle: (event, data) => {
|
||||
aggregatorHandlers.on_chat_model_end.handle(event, data);
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
collectedUsage.push(usage);
|
||||
}
|
||||
},
|
||||
},
|
||||
on_tool_end: new ToolEndHandler(toolEndCallback, logger),
|
||||
on_run_step_completed: { handle: () => {} },
|
||||
on_chain_stream: { handle: () => {} },
|
||||
on_chain_end: { handle: () => {} },
|
||||
on_agent_update: { handle: () => {} },
|
||||
on_custom_event: { handle: () => {} },
|
||||
on_tool_execute: createToolExecuteHandler(toolExecuteOptions),
|
||||
};
|
||||
|
||||
const userId = req.user?.id ?? 'api-user';
|
||||
const userMCPAuthMap = primaryConfig.userMCPAuthMap;
|
||||
|
||||
const run = await createRun({
|
||||
agents: [primaryConfig],
|
||||
messages: formattedMessages,
|
||||
indexTokenCountMap,
|
||||
runId: responseId,
|
||||
signal: abortController.signal,
|
||||
customHandlers: handlers,
|
||||
requestBody: {
|
||||
messageId: responseId,
|
||||
conversationId,
|
||||
},
|
||||
user: { id: userId },
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
throw new Error('Failed to create agent run');
|
||||
}
|
||||
|
||||
const config = {
|
||||
runName: 'AgentRun',
|
||||
configurable: {
|
||||
thread_id: conversationId,
|
||||
user_id: userId,
|
||||
user: createSafeUser(req.user),
|
||||
...(userMCPAuthMap != null && { userMCPAuthMap }),
|
||||
},
|
||||
signal: abortController.signal,
|
||||
streamMode: 'values',
|
||||
version: 'v2',
|
||||
};
|
||||
|
||||
await run.processStream({ messages: formattedMessages }, config, {
|
||||
callbacks: {
|
||||
[Callback.TOOL_ERROR]: (graph, error, toolId) => {
|
||||
logger.error(`[Responses API] Tool Error "${toolId}"`, error);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Record token usage against balance
|
||||
const balanceConfig = getBalanceConfig(req.config);
|
||||
const transactionsConfig = getTransactionsConfig(req.config);
|
||||
recordCollectedUsage(
|
||||
{ spendTokens, spendStructuredTokens },
|
||||
{
|
||||
user: userId,
|
||||
conversationId,
|
||||
collectedUsage,
|
||||
context: 'message',
|
||||
balance: balanceConfig,
|
||||
transactions: transactionsConfig,
|
||||
model: primaryConfig.model || agent.model_parameters?.model,
|
||||
},
|
||||
).catch((err) => {
|
||||
logger.error('[Responses API] Error recording usage:', err);
|
||||
});
|
||||
|
||||
if (artifactPromises.length > 0) {
|
||||
try {
|
||||
await Promise.all(artifactPromises);
|
||||
} catch (artifactError) {
|
||||
logger.warn('[Responses API] Error processing artifacts:', artifactError);
|
||||
}
|
||||
}
|
||||
|
||||
const response = buildAggregatedResponse(context, aggregator);
|
||||
|
||||
if (request.store === true) {
|
||||
try {
|
||||
await saveConversation(req, conversationId, agentId, agent);
|
||||
|
||||
await saveInputMessages(req, conversationId, inputMessages, agentId);
|
||||
|
||||
await saveResponseOutput(req, conversationId, responseId, response, agentId);
|
||||
|
||||
logger.debug(
|
||||
`[Responses API] Stored response ${responseId} in conversation ${conversationId}`,
|
||||
);
|
||||
} catch (saveError) {
|
||||
logger.error('[Responses API] Error saving response:', saveError);
|
||||
// Don't fail the request if saving fails
|
||||
}
|
||||
}
|
||||
|
||||
res.json(response);
|
||||
|
||||
const duration = Date.now() - requestStartTime;
|
||||
logger.debug(
|
||||
`[Responses API] Request ${responseId} completed in ${duration}ms (non-streaming)`,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'An error occurred';
|
||||
logger.error('[Responses API] Error:', error);
|
||||
|
||||
// Check if we already started streaming (headers sent)
|
||||
if (res.headersSent) {
|
||||
// Headers already sent, write error event and close
|
||||
writeDone(res);
|
||||
res.end();
|
||||
} else {
|
||||
// Forward upstream provider status codes (e.g., Anthropic 400s) instead of masking as 500
|
||||
const statusCode =
|
||||
typeof error?.status === 'number' && error.status >= 400 && error.status < 600
|
||||
? error.status
|
||||
: 500;
|
||||
const errorType = statusCode >= 400 && statusCode < 500 ? 'invalid_request' : 'server_error';
|
||||
sendResponsesErrorResponse(res, statusCode, errorMessage, errorType);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* List available agents as models - GET /v1/models (also works with /v1/responses/models)
|
||||
*
|
||||
* Returns a list of available agents the user has remote access to.
|
||||
*
|
||||
* @param {import('express').Request} req
|
||||
* @param {import('express').Response} res
|
||||
*/
|
||||
const listModels = async (req, res) => {
|
||||
try {
|
||||
const userId = req.user?.id;
|
||||
const userRole = req.user?.role;
|
||||
|
||||
if (!userId) {
|
||||
return sendResponsesErrorResponse(res, 401, 'Authentication required', 'auth_error');
|
||||
}
|
||||
|
||||
// Find agents the user has remote access to (VIEW permission on REMOTE_AGENT)
|
||||
const accessibleAgentIds = await findAccessibleResources({
|
||||
userId,
|
||||
role: userRole,
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
requiredPermissions: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
// Get the accessible agents
|
||||
let agents = [];
|
||||
if (accessibleAgentIds.length > 0) {
|
||||
agents = await getAgents({ _id: { $in: accessibleAgentIds } });
|
||||
}
|
||||
|
||||
// Convert to models format
|
||||
const models = agents.map((agent) => ({
|
||||
id: agent.id,
|
||||
object: 'model',
|
||||
created: Math.floor(new Date(agent.createdAt).getTime() / 1000),
|
||||
owned_by: agent.author ?? 'librechat',
|
||||
// Additional metadata
|
||||
name: agent.name,
|
||||
description: agent.description,
|
||||
provider: agent.provider,
|
||||
}));
|
||||
|
||||
res.json({
|
||||
object: 'list',
|
||||
data: models,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[Responses API] Error listing models:', error);
|
||||
sendResponsesErrorResponse(
|
||||
res,
|
||||
500,
|
||||
error instanceof Error ? error.message : 'Failed to list models',
|
||||
'server_error',
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Get Response - GET /v1/responses/:id
|
||||
*
|
||||
* Retrieves a stored response by its ID.
|
||||
* The response ID maps to a conversationId in LibreChat's storage.
|
||||
*
|
||||
* @param {import('express').Request} req
|
||||
* @param {import('express').Response} res
|
||||
*/
|
||||
const getResponse = async (req, res) => {
|
||||
try {
|
||||
const responseId = req.params.id;
|
||||
const userId = req.user?.id;
|
||||
|
||||
if (!responseId) {
|
||||
return sendResponsesErrorResponse(res, 400, 'Response ID is required');
|
||||
}
|
||||
|
||||
// The responseId could be either the response ID or the conversation ID
|
||||
// Try to find a conversation with this ID
|
||||
const conversation = await getConvo(userId, responseId);
|
||||
|
||||
if (!conversation) {
|
||||
return sendResponsesErrorResponse(
|
||||
res,
|
||||
404,
|
||||
`Response not found: ${responseId}`,
|
||||
'not_found',
|
||||
'response_not_found',
|
||||
);
|
||||
}
|
||||
|
||||
// Load messages for this conversation
|
||||
const messages = await db.getMessages({ conversationId: responseId, user: userId });
|
||||
|
||||
if (!messages || messages.length === 0) {
|
||||
return sendResponsesErrorResponse(
|
||||
res,
|
||||
404,
|
||||
`No messages found for response: ${responseId}`,
|
||||
'not_found',
|
||||
'response_not_found',
|
||||
);
|
||||
}
|
||||
|
||||
// Convert messages to Open Responses output format
|
||||
const output = convertMessagesToOutputItems(messages);
|
||||
|
||||
// Find the last assistant message for usage info
|
||||
const lastAssistantMessage = messages.filter((m) => !m.isCreatedByUser).pop();
|
||||
|
||||
// Build the response object
|
||||
const response = {
|
||||
id: responseId,
|
||||
object: 'response',
|
||||
created_at: Math.floor(new Date(conversation.createdAt || Date.now()).getTime() / 1000),
|
||||
completed_at: Math.floor(new Date(conversation.updatedAt || Date.now()).getTime() / 1000),
|
||||
status: 'completed',
|
||||
incomplete_details: null,
|
||||
model: conversation.agentId || conversation.model || 'unknown',
|
||||
previous_response_id: null,
|
||||
instructions: null,
|
||||
output,
|
||||
error: null,
|
||||
tools: [],
|
||||
tool_choice: 'auto',
|
||||
truncation: 'disabled',
|
||||
parallel_tool_calls: true,
|
||||
text: { format: { type: 'text' } },
|
||||
temperature: 1,
|
||||
top_p: 1,
|
||||
presence_penalty: 0,
|
||||
frequency_penalty: 0,
|
||||
top_logprobs: null,
|
||||
reasoning: null,
|
||||
user: userId,
|
||||
usage: lastAssistantMessage?.tokenCount
|
||||
? {
|
||||
input_tokens: 0,
|
||||
output_tokens: lastAssistantMessage.tokenCount,
|
||||
total_tokens: lastAssistantMessage.tokenCount,
|
||||
}
|
||||
: null,
|
||||
max_output_tokens: null,
|
||||
max_tool_calls: null,
|
||||
store: true,
|
||||
background: false,
|
||||
service_tier: 'default',
|
||||
metadata: {},
|
||||
safety_identifier: null,
|
||||
prompt_cache_key: null,
|
||||
};
|
||||
|
||||
res.json(response);
|
||||
} catch (error) {
|
||||
logger.error('[Responses API] Error getting response:', error);
|
||||
sendResponsesErrorResponse(
|
||||
res,
|
||||
500,
|
||||
error instanceof Error ? error.message : 'Failed to get response',
|
||||
'server_error',
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
createResponse,
|
||||
getResponse,
|
||||
listModels,
|
||||
setAppConfig,
|
||||
};
|
||||
|
|
@ -11,7 +11,9 @@ const {
|
|||
convertOcrToContextInPlace,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
Time,
|
||||
Tools,
|
||||
CacheKeys,
|
||||
Constants,
|
||||
FileSources,
|
||||
ResourceType,
|
||||
|
|
@ -21,8 +23,6 @@ const {
|
|||
PermissionBits,
|
||||
actionDelimiter,
|
||||
removeNullishValues,
|
||||
CacheKeys,
|
||||
Time,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
getListAgentsByAccess,
|
||||
|
|
@ -94,16 +94,25 @@ const createAgentHandler = async (req, res) => {
|
|||
|
||||
const agent = await createAgent(agentData);
|
||||
|
||||
// Automatically grant owner permissions to the creator
|
||||
try {
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
});
|
||||
await Promise.all([
|
||||
grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
}),
|
||||
grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.REMOTE_AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
}),
|
||||
]);
|
||||
logger.debug(
|
||||
`[createAgent] Granted owner permissions to user ${userId} for agent ${agent.id}`,
|
||||
);
|
||||
|
|
@ -396,16 +405,25 @@ const duplicateAgentHandler = async (req, res) => {
|
|||
newAgentData.actions = agentActions;
|
||||
const newAgent = await createAgent(newAgentData);
|
||||
|
||||
// Automatically grant owner permissions to the duplicator
|
||||
try {
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: newAgent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
});
|
||||
await Promise.all([
|
||||
grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: newAgent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
}),
|
||||
grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
resourceId: newAgent._id,
|
||||
accessRoleId: AccessRoleIds.REMOTE_AGENT_OWNER,
|
||||
grantedBy: userId,
|
||||
}),
|
||||
]);
|
||||
logger.debug(
|
||||
`[duplicateAgent] Granted owner permissions to user ${userId} for duplicated agent ${newAgent.id}`,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ const logoutController = async (req, res) => {
|
|||
|
||||
res.clearCookie('refreshToken');
|
||||
res.clearCookie('openid_access_token');
|
||||
res.clearCookie('openid_id_token');
|
||||
res.clearCookie('openid_user_id');
|
||||
res.clearCookie('token_provider');
|
||||
const response = { message };
|
||||
|
|
|
|||
79
api/server/controllers/auth/oauth.js
Normal file
79
api/server/controllers/auth/oauth.js
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const { logger, DEFAULT_SESSION_EXPIRY } = require('@librechat/data-schemas');
|
||||
const {
|
||||
isEnabled,
|
||||
getAdminPanelUrl,
|
||||
isAdminPanelRedirect,
|
||||
generateAdminExchangeCode,
|
||||
} = require('@librechat/api');
|
||||
const { syncUserEntraGroupMemberships } = require('~/server/services/PermissionService');
|
||||
const { setAuthTokens, setOpenIDAuthTokens } = require('~/server/services/AuthService');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
const { checkBan } = require('~/server/middleware');
|
||||
const { generateToken } = require('~/models');
|
||||
|
||||
const domains = {
|
||||
client: process.env.DOMAIN_CLIENT,
|
||||
server: process.env.DOMAIN_SERVER,
|
||||
};
|
||||
|
||||
function createOAuthHandler(redirectUri = domains.client) {
|
||||
/**
|
||||
* A handler to process OAuth authentication results.
|
||||
* @type {Function}
|
||||
* @param {ServerRequest} req - Express request object.
|
||||
* @param {ServerResponse} res - Express response object.
|
||||
* @param {NextFunction} next - Express next middleware function.
|
||||
*/
|
||||
return async (req, res, next) => {
|
||||
try {
|
||||
if (res.headersSent) {
|
||||
return;
|
||||
}
|
||||
|
||||
await checkBan(req, res);
|
||||
if (req.banned) {
|
||||
return;
|
||||
}
|
||||
|
||||
/** Check if this is an admin panel redirect (cross-origin) */
|
||||
if (isAdminPanelRedirect(redirectUri, getAdminPanelUrl(), domains.client)) {
|
||||
/** For admin panel, generate exchange code instead of setting cookies */
|
||||
const cache = getLogStores(CacheKeys.ADMIN_OAUTH_EXCHANGE);
|
||||
const sessionExpiry = Number(process.env.SESSION_EXPIRY) || DEFAULT_SESSION_EXPIRY;
|
||||
const token = await generateToken(req.user, sessionExpiry);
|
||||
|
||||
/** Get refresh token from tokenset for OpenID users */
|
||||
const refreshToken =
|
||||
req.user.tokenset?.refresh_token || req.user.federatedTokens?.refresh_token;
|
||||
|
||||
const exchangeCode = await generateAdminExchangeCode(cache, req.user, token, refreshToken);
|
||||
|
||||
const callbackUrl = new URL(redirectUri);
|
||||
callbackUrl.searchParams.set('code', exchangeCode);
|
||||
logger.info(`[OAuth] Admin panel redirect with exchange code for user: ${req.user.email}`);
|
||||
return res.redirect(callbackUrl.toString());
|
||||
}
|
||||
|
||||
/** Standard OAuth flow - set cookies and redirect */
|
||||
if (
|
||||
req.user &&
|
||||
req.user.provider == 'openid' &&
|
||||
isEnabled(process.env.OPENID_REUSE_TOKENS) === true
|
||||
) {
|
||||
await syncUserEntraGroupMemberships(req.user, req.user.tokenset.access_token);
|
||||
setOpenIDAuthTokens(req.user.tokenset, req, res, req.user._id.toString());
|
||||
} else {
|
||||
await setAuthTokens(req.user._id, res);
|
||||
}
|
||||
res.redirect(redirectUri);
|
||||
} catch (err) {
|
||||
logger.error('Error in setting authentication tokens:', err);
|
||||
next(err);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
createOAuthHandler,
|
||||
};
|
||||
|
|
@ -299,6 +299,7 @@ if (cluster.isMaster) {
|
|||
app.use('/api/auth', routes.auth);
|
||||
app.use('/api/actions', routes.actions);
|
||||
app.use('/api/keys', routes.keys);
|
||||
app.use('/api/api-keys', routes.apiKeys);
|
||||
app.use('/api/user', routes.user);
|
||||
app.use('/api/search', routes.search);
|
||||
app.use('/api/messages', routes.messages);
|
||||
|
|
|
|||
|
|
@ -134,8 +134,10 @@ const startServer = async () => {
|
|||
app.use('/oauth', routes.oauth);
|
||||
/* API Endpoints */
|
||||
app.use('/api/auth', routes.auth);
|
||||
app.use('/api/admin', routes.adminAuth);
|
||||
app.use('/api/actions', routes.actions);
|
||||
app.use('/api/keys', routes.keys);
|
||||
app.use('/api/api-keys', routes.apiKeys);
|
||||
app.use('/api/user', routes.user);
|
||||
app.use('/api/search', routes.search);
|
||||
app.use('/api/messages', routes.messages);
|
||||
|
|
@ -249,6 +251,15 @@ process.on('uncaughtException', (err) => {
|
|||
return;
|
||||
}
|
||||
|
||||
if (isEnabled(process.env.CONTINUE_ON_UNCAUGHT_EXCEPTION)) {
|
||||
logger.error('Unhandled error encountered. The app will continue running.', {
|
||||
name: err?.name,
|
||||
message: err?.message,
|
||||
stack: err?.stack,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -7,13 +7,89 @@ const {
|
|||
sanitizeMessageForTransmit,
|
||||
} = require('@librechat/api');
|
||||
const { isAssistantsEndpoint, ErrorTypes } = require('librechat-data-provider');
|
||||
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
|
||||
const { truncateText, smartTruncateText } = require('~/app/clients/prompts');
|
||||
const clearPendingReq = require('~/cache/clearPendingReq');
|
||||
const { sendError } = require('~/server/middleware/error');
|
||||
const { spendTokens } = require('~/models/spendTokens');
|
||||
const { saveMessage, getConvo } = require('~/models');
|
||||
const { abortRun } = require('./abortRun');
|
||||
|
||||
/**
|
||||
* Spend tokens for all models from collected usage.
|
||||
* This handles both sequential and parallel agent execution.
|
||||
*
|
||||
* IMPORTANT: After spending, this function clears the collectedUsage array
|
||||
* to prevent double-spending. The array is shared with AgentClient.collectedUsage,
|
||||
* so clearing it here prevents the finally block from also spending tokens.
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {string} params.userId - User ID
|
||||
* @param {string} params.conversationId - Conversation ID
|
||||
* @param {Array<Object>} params.collectedUsage - Usage metadata from all models
|
||||
* @param {string} [params.fallbackModel] - Fallback model name if not in usage
|
||||
*/
|
||||
async function spendCollectedUsage({ userId, conversationId, collectedUsage, fallbackModel }) {
|
||||
if (!collectedUsage || collectedUsage.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const spendPromises = [];
|
||||
|
||||
for (const usage of collectedUsage) {
|
||||
if (!usage) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Support both OpenAI format (input_token_details) and Anthropic format (cache_*_input_tokens)
|
||||
const cache_creation =
|
||||
Number(usage.input_token_details?.cache_creation) ||
|
||||
Number(usage.cache_creation_input_tokens) ||
|
||||
0;
|
||||
const cache_read =
|
||||
Number(usage.input_token_details?.cache_read) || Number(usage.cache_read_input_tokens) || 0;
|
||||
|
||||
const txMetadata = {
|
||||
context: 'abort',
|
||||
conversationId,
|
||||
user: userId,
|
||||
model: usage.model ?? fallbackModel,
|
||||
};
|
||||
|
||||
if (cache_creation > 0 || cache_read > 0) {
|
||||
spendPromises.push(
|
||||
spendStructuredTokens(txMetadata, {
|
||||
promptTokens: {
|
||||
input: usage.input_tokens,
|
||||
write: cache_creation,
|
||||
read: cache_read,
|
||||
},
|
||||
completionTokens: usage.output_tokens,
|
||||
}).catch((err) => {
|
||||
logger.error('[abortMiddleware] Error spending structured tokens for abort', err);
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
spendPromises.push(
|
||||
spendTokens(txMetadata, {
|
||||
promptTokens: usage.input_tokens,
|
||||
completionTokens: usage.output_tokens,
|
||||
}).catch((err) => {
|
||||
logger.error('[abortMiddleware] Error spending tokens for abort', err);
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Wait for all token spending to complete
|
||||
await Promise.all(spendPromises);
|
||||
|
||||
// Clear the array to prevent double-spending from the AgentClient finally block.
|
||||
// The collectedUsage array is shared by reference with AgentClient.collectedUsage,
|
||||
// so clearing it here ensures recordCollectedUsage() sees an empty array and returns early.
|
||||
collectedUsage.length = 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Abort an active message generation.
|
||||
* Uses GenerationJobManager for all agent requests.
|
||||
|
|
@ -39,9 +115,8 @@ async function abortMessage(req, res) {
|
|||
return;
|
||||
}
|
||||
|
||||
const { jobData, content, text } = abortResult;
|
||||
const { jobData, content, text, collectedUsage } = abortResult;
|
||||
|
||||
// Count tokens and spend them
|
||||
const completionTokens = await countTokens(text);
|
||||
const promptTokens = jobData?.promptTokens ?? 0;
|
||||
|
||||
|
|
@ -62,10 +137,21 @@ async function abortMessage(req, res) {
|
|||
tokenCount: completionTokens,
|
||||
};
|
||||
|
||||
await spendTokens(
|
||||
{ ...responseMessage, context: 'incomplete', user: userId },
|
||||
{ promptTokens, completionTokens },
|
||||
);
|
||||
// Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo)
|
||||
if (collectedUsage && collectedUsage.length > 0) {
|
||||
await spendCollectedUsage({
|
||||
userId,
|
||||
conversationId: jobData?.conversationId,
|
||||
collectedUsage,
|
||||
fallbackModel: jobData?.model,
|
||||
});
|
||||
} else {
|
||||
// Fallback: no collected usage, use text-based token counting for primary model only
|
||||
await spendTokens(
|
||||
{ ...responseMessage, context: 'incomplete', user: userId },
|
||||
{ promptTokens, completionTokens },
|
||||
);
|
||||
}
|
||||
|
||||
await saveMessage(
|
||||
req,
|
||||
|
|
|
|||
428
api/server/middleware/abortMiddleware.spec.js
Normal file
428
api/server/middleware/abortMiddleware.spec.js
Normal file
|
|
@ -0,0 +1,428 @@
|
|||
/**
|
||||
* Tests for abortMiddleware - spendCollectedUsage function
|
||||
*
|
||||
* This tests the token spending logic for abort scenarios,
|
||||
* particularly for parallel agents (addedConvo) where multiple
|
||||
* models need their tokens spent.
|
||||
*/
|
||||
|
||||
const mockSpendTokens = jest.fn().mockResolvedValue();
|
||||
const mockSpendStructuredTokens = jest.fn().mockResolvedValue();
|
||||
|
||||
jest.mock('~/models/spendTokens', () => ({
|
||||
spendTokens: (...args) => mockSpendTokens(...args),
|
||||
spendStructuredTokens: (...args) => mockSpendStructuredTokens(...args),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
info: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
countTokens: jest.fn().mockResolvedValue(100),
|
||||
isEnabled: jest.fn().mockReturnValue(false),
|
||||
sendEvent: jest.fn(),
|
||||
GenerationJobManager: {
|
||||
abortJob: jest.fn(),
|
||||
},
|
||||
sanitizeMessageForTransmit: jest.fn((msg) => msg),
|
||||
}));
|
||||
|
||||
jest.mock('librechat-data-provider', () => ({
|
||||
isAssistantsEndpoint: jest.fn().mockReturnValue(false),
|
||||
ErrorTypes: { INVALID_REQUEST: 'INVALID_REQUEST', NO_SYSTEM_MESSAGES: 'NO_SYSTEM_MESSAGES' },
|
||||
}));
|
||||
|
||||
jest.mock('~/app/clients/prompts', () => ({
|
||||
truncateText: jest.fn((text) => text),
|
||||
smartTruncateText: jest.fn((text) => text),
|
||||
}));
|
||||
|
||||
jest.mock('~/cache/clearPendingReq', () => jest.fn().mockResolvedValue());
|
||||
|
||||
jest.mock('~/server/middleware/error', () => ({
|
||||
sendError: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: jest.fn().mockResolvedValue(),
|
||||
getConvo: jest.fn().mockResolvedValue({ title: 'Test Chat' }),
|
||||
}));
|
||||
|
||||
jest.mock('./abortRun', () => ({
|
||||
abortRun: jest.fn(),
|
||||
}));
|
||||
|
||||
// Import the module after mocks are set up
|
||||
// We need to extract the spendCollectedUsage function for testing
|
||||
// Since it's not exported, we'll test it through the handleAbort flow
|
||||
|
||||
describe('abortMiddleware - spendCollectedUsage', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('spendCollectedUsage logic', () => {
|
||||
// Since spendCollectedUsage is not exported, we test the logic directly
|
||||
// by replicating the function here for unit testing
|
||||
|
||||
const spendCollectedUsage = async ({
|
||||
userId,
|
||||
conversationId,
|
||||
collectedUsage,
|
||||
fallbackModel,
|
||||
}) => {
|
||||
if (!collectedUsage || collectedUsage.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const spendPromises = [];
|
||||
|
||||
for (const usage of collectedUsage) {
|
||||
if (!usage) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const cache_creation =
|
||||
Number(usage.input_token_details?.cache_creation) ||
|
||||
Number(usage.cache_creation_input_tokens) ||
|
||||
0;
|
||||
const cache_read =
|
||||
Number(usage.input_token_details?.cache_read) ||
|
||||
Number(usage.cache_read_input_tokens) ||
|
||||
0;
|
||||
|
||||
const txMetadata = {
|
||||
context: 'abort',
|
||||
conversationId,
|
||||
user: userId,
|
||||
model: usage.model ?? fallbackModel,
|
||||
};
|
||||
|
||||
if (cache_creation > 0 || cache_read > 0) {
|
||||
spendPromises.push(
|
||||
mockSpendStructuredTokens(txMetadata, {
|
||||
promptTokens: {
|
||||
input: usage.input_tokens,
|
||||
write: cache_creation,
|
||||
read: cache_read,
|
||||
},
|
||||
completionTokens: usage.output_tokens,
|
||||
}).catch(() => {
|
||||
// Log error but don't throw
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
spendPromises.push(
|
||||
mockSpendTokens(txMetadata, {
|
||||
promptTokens: usage.input_tokens,
|
||||
completionTokens: usage.output_tokens,
|
||||
}).catch(() => {
|
||||
// Log error but don't throw
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Wait for all token spending to complete
|
||||
await Promise.all(spendPromises);
|
||||
|
||||
// Clear the array to prevent double-spending
|
||||
collectedUsage.length = 0;
|
||||
};
|
||||
|
||||
it('should return early if collectedUsage is empty', async () => {
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage: [],
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).not.toHaveBeenCalled();
|
||||
expect(mockSpendStructuredTokens).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return early if collectedUsage is null', async () => {
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage: null,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).not.toHaveBeenCalled();
|
||||
expect(mockSpendStructuredTokens).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip null entries in collectedUsage', async () => {
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' },
|
||||
null,
|
||||
{ input_tokens: 200, output_tokens: 60, model: 'gpt-4' },
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should spend tokens for single model', async () => {
|
||||
const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(1);
|
||||
expect(mockSpendTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
context: 'abort',
|
||||
conversationId: 'convo-123',
|
||||
user: 'user-123',
|
||||
model: 'gpt-4',
|
||||
}),
|
||||
{ promptTokens: 100, completionTokens: 50 },
|
||||
);
|
||||
});
|
||||
|
||||
it('should spend tokens for multiple models (parallel agents)', async () => {
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' },
|
||||
{ input_tokens: 80, output_tokens: 40, model: 'claude-3' },
|
||||
{ input_tokens: 120, output_tokens: 60, model: 'gemini-pro' },
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(3);
|
||||
|
||||
// Verify each model was called
|
||||
expect(mockSpendTokens).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
expect.objectContaining({ model: 'gpt-4' }),
|
||||
{ promptTokens: 100, completionTokens: 50 },
|
||||
);
|
||||
expect(mockSpendTokens).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
expect.objectContaining({ model: 'claude-3' }),
|
||||
{ promptTokens: 80, completionTokens: 40 },
|
||||
);
|
||||
expect(mockSpendTokens).toHaveBeenNthCalledWith(
|
||||
3,
|
||||
expect.objectContaining({ model: 'gemini-pro' }),
|
||||
{ promptTokens: 120, completionTokens: 60 },
|
||||
);
|
||||
});
|
||||
|
||||
it('should use fallbackModel when usage.model is missing', async () => {
|
||||
const collectedUsage = [{ input_tokens: 100, output_tokens: 50 }];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'fallback-model',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ model: 'fallback-model' }),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it('should use spendStructuredTokens for OpenAI format cache tokens', async () => {
|
||||
const collectedUsage = [
|
||||
{
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
model: 'gpt-4',
|
||||
input_token_details: {
|
||||
cache_creation: 20,
|
||||
cache_read: 10,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendStructuredTokens).toHaveBeenCalledTimes(1);
|
||||
expect(mockSpendTokens).not.toHaveBeenCalled();
|
||||
expect(mockSpendStructuredTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ model: 'gpt-4', context: 'abort' }),
|
||||
{
|
||||
promptTokens: {
|
||||
input: 100,
|
||||
write: 20,
|
||||
read: 10,
|
||||
},
|
||||
completionTokens: 50,
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it('should use spendStructuredTokens for Anthropic format cache tokens', async () => {
|
||||
const collectedUsage = [
|
||||
{
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
model: 'claude-3',
|
||||
cache_creation_input_tokens: 25,
|
||||
cache_read_input_tokens: 15,
|
||||
},
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'claude-3',
|
||||
});
|
||||
|
||||
expect(mockSpendStructuredTokens).toHaveBeenCalledTimes(1);
|
||||
expect(mockSpendTokens).not.toHaveBeenCalled();
|
||||
expect(mockSpendStructuredTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ model: 'claude-3' }),
|
||||
{
|
||||
promptTokens: {
|
||||
input: 100,
|
||||
write: 25,
|
||||
read: 15,
|
||||
},
|
||||
completionTokens: 50,
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle mixed cache and non-cache entries', async () => {
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' },
|
||||
{
|
||||
input_tokens: 150,
|
||||
output_tokens: 30,
|
||||
model: 'claude-3',
|
||||
cache_creation_input_tokens: 20,
|
||||
cache_read_input_tokens: 10,
|
||||
},
|
||||
{ input_tokens: 200, output_tokens: 20, model: 'gemini-pro' },
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(2);
|
||||
expect(mockSpendStructuredTokens).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle real-world parallel agent abort scenario', async () => {
|
||||
// Simulates: Primary agent (gemini) + addedConvo agent (gpt-5) aborted mid-stream
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 31596, output_tokens: 151, model: 'gemini-3-flash-preview' },
|
||||
{ input_tokens: 28000, output_tokens: 120, model: 'gpt-5.2' },
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gemini-3-flash-preview',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(2);
|
||||
|
||||
// Primary model
|
||||
expect(mockSpendTokens).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
expect.objectContaining({ model: 'gemini-3-flash-preview' }),
|
||||
{ promptTokens: 31596, completionTokens: 151 },
|
||||
);
|
||||
|
||||
// Parallel model (addedConvo)
|
||||
expect(mockSpendTokens).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
expect.objectContaining({ model: 'gpt-5.2' }),
|
||||
{ promptTokens: 28000, completionTokens: 120 },
|
||||
);
|
||||
});
|
||||
|
||||
it('should clear collectedUsage array after spending to prevent double-spending', async () => {
|
||||
// This tests the race condition fix: after abort middleware spends tokens,
|
||||
// the collectedUsage array is cleared so AgentClient.recordCollectedUsage()
|
||||
// (which shares the same array reference) sees an empty array and returns early.
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' },
|
||||
{ input_tokens: 80, output_tokens: 40, model: 'claude-3' },
|
||||
];
|
||||
|
||||
expect(collectedUsage.length).toBe(2);
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(2);
|
||||
|
||||
// The array should be cleared after spending
|
||||
expect(collectedUsage.length).toBe(0);
|
||||
});
|
||||
|
||||
it('should await all token spending operations before clearing array', async () => {
|
||||
// Ensure we don't clear the array before spending completes
|
||||
let spendCallCount = 0;
|
||||
mockSpendTokens.mockImplementation(async () => {
|
||||
spendCallCount++;
|
||||
// Simulate async delay
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
});
|
||||
|
||||
const collectedUsage = [
|
||||
{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' },
|
||||
{ input_tokens: 80, output_tokens: 40, model: 'claude-3' },
|
||||
];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
});
|
||||
|
||||
// Both spend calls should have completed
|
||||
expect(spendCallCount).toBe(2);
|
||||
|
||||
// Array should be cleared after awaiting
|
||||
expect(collectedUsage.length).toBe(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -5,9 +5,11 @@ const {
|
|||
EModelEndpoint,
|
||||
isAgentsEndpoint,
|
||||
parseCompactConvo,
|
||||
getDefaultParamsEndpoint,
|
||||
} = require('librechat-data-provider');
|
||||
const azureAssistants = require('~/server/services/Endpoints/azureAssistants');
|
||||
const assistants = require('~/server/services/Endpoints/assistants');
|
||||
const { getEndpointsConfig } = require('~/server/services/Config');
|
||||
const agents = require('~/server/services/Endpoints/agents');
|
||||
const { updateFilesUsage } = require('~/models');
|
||||
|
||||
|
|
@ -19,9 +21,24 @@ const buildFunction = {
|
|||
|
||||
async function buildEndpointOption(req, res, next) {
|
||||
const { endpoint, endpointType } = req.body;
|
||||
|
||||
let endpointsConfig;
|
||||
try {
|
||||
endpointsConfig = await getEndpointsConfig(req);
|
||||
} catch (error) {
|
||||
logger.error('Error fetching endpoints config in buildEndpointOption', error);
|
||||
}
|
||||
|
||||
const defaultParamsEndpoint = getDefaultParamsEndpoint(endpointsConfig, endpoint);
|
||||
|
||||
let parsedBody;
|
||||
try {
|
||||
parsedBody = parseCompactConvo({ endpoint, endpointType, conversation: req.body });
|
||||
parsedBody = parseCompactConvo({
|
||||
endpoint,
|
||||
endpointType,
|
||||
conversation: req.body,
|
||||
defaultParamsEndpoint,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error(`Error parsing compact conversation for endpoint ${endpoint}`, error);
|
||||
logger.debug({
|
||||
|
|
@ -55,6 +72,7 @@ async function buildEndpointOption(req, res, next) {
|
|||
endpoint,
|
||||
endpointType,
|
||||
conversation: currentModelSpec.preset,
|
||||
defaultParamsEndpoint,
|
||||
});
|
||||
if (currentModelSpec.iconURL != null && currentModelSpec.iconURL !== '') {
|
||||
parsedBody.iconURL = currentModelSpec.iconURL;
|
||||
|
|
|
|||
237
api/server/middleware/buildEndpointOption.spec.js
Normal file
237
api/server/middleware/buildEndpointOption.spec.js
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
/**
|
||||
* Wrap parseCompactConvo: the REAL function runs, but jest can observe
|
||||
* calls and return values. Must be declared before require('./buildEndpointOption')
|
||||
* so the destructured reference in the middleware captures the wrapper.
|
||||
*/
|
||||
jest.mock('librechat-data-provider', () => {
|
||||
const actual = jest.requireActual('librechat-data-provider');
|
||||
return {
|
||||
...actual,
|
||||
parseCompactConvo: jest.fn((...args) => actual.parseCompactConvo(...args)),
|
||||
};
|
||||
});
|
||||
|
||||
const { EModelEndpoint, parseCompactConvo } = require('librechat-data-provider');
|
||||
|
||||
const mockBuildOptions = jest.fn((_endpoint, parsedBody) => ({
|
||||
...parsedBody,
|
||||
endpoint: _endpoint,
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Endpoints/azureAssistants', () => ({
|
||||
buildOptions: mockBuildOptions,
|
||||
}));
|
||||
jest.mock('~/server/services/Endpoints/assistants', () => ({
|
||||
buildOptions: mockBuildOptions,
|
||||
}));
|
||||
jest.mock('~/server/services/Endpoints/agents', () => ({
|
||||
buildOptions: mockBuildOptions,
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
updateFilesUsage: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockGetEndpointsConfig = jest.fn();
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
getEndpointsConfig: (...args) => mockGetEndpointsConfig(...args),
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
handleError: jest.fn(),
|
||||
}));
|
||||
|
||||
const buildEndpointOption = require('./buildEndpointOption');
|
||||
|
||||
const createReq = (body, config = {}) => ({
|
||||
body,
|
||||
config,
|
||||
baseUrl: '/api/chat',
|
||||
});
|
||||
|
||||
const createRes = () => ({
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
});
|
||||
|
||||
describe('buildEndpointOption - defaultParamsEndpoint parsing', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should pass defaultParamsEndpoint to parseCompactConvo and preserve maxOutputTokens', async () => {
|
||||
mockGetEndpointsConfig.mockResolvedValue({
|
||||
AnthropicClaude: {
|
||||
type: EModelEndpoint.custom,
|
||||
customParams: {
|
||||
defaultParamsEndpoint: EModelEndpoint.anthropic,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const req = createReq(
|
||||
{
|
||||
endpoint: 'AnthropicClaude',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
model: 'anthropic/claude-opus-4.5',
|
||||
temperature: 0.7,
|
||||
maxOutputTokens: 8192,
|
||||
topP: 0.9,
|
||||
maxContextTokens: 50000,
|
||||
},
|
||||
{ modelSpecs: null },
|
||||
);
|
||||
|
||||
await buildEndpointOption(req, createRes(), jest.fn());
|
||||
|
||||
expect(parseCompactConvo).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
defaultParamsEndpoint: EModelEndpoint.anthropic,
|
||||
}),
|
||||
);
|
||||
|
||||
const parsedResult = parseCompactConvo.mock.results[0].value;
|
||||
expect(parsedResult.maxOutputTokens).toBe(8192);
|
||||
expect(parsedResult.topP).toBe(0.9);
|
||||
expect(parsedResult.temperature).toBe(0.7);
|
||||
expect(parsedResult.maxContextTokens).toBe(50000);
|
||||
});
|
||||
|
||||
it('should strip maxOutputTokens when no defaultParamsEndpoint is configured', async () => {
|
||||
mockGetEndpointsConfig.mockResolvedValue({
|
||||
MyOpenRouter: {
|
||||
type: EModelEndpoint.custom,
|
||||
},
|
||||
});
|
||||
|
||||
const req = createReq(
|
||||
{
|
||||
endpoint: 'MyOpenRouter',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
model: 'gpt-4o',
|
||||
temperature: 0.7,
|
||||
maxOutputTokens: 8192,
|
||||
max_tokens: 4096,
|
||||
},
|
||||
{ modelSpecs: null },
|
||||
);
|
||||
|
||||
await buildEndpointOption(req, createRes(), jest.fn());
|
||||
|
||||
expect(parseCompactConvo).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
defaultParamsEndpoint: undefined,
|
||||
}),
|
||||
);
|
||||
|
||||
const parsedResult = parseCompactConvo.mock.results[0].value;
|
||||
expect(parsedResult.maxOutputTokens).toBeUndefined();
|
||||
expect(parsedResult.max_tokens).toBe(4096);
|
||||
expect(parsedResult.temperature).toBe(0.7);
|
||||
});
|
||||
|
||||
it('should strip bedrock region from custom endpoint without defaultParamsEndpoint', async () => {
|
||||
mockGetEndpointsConfig.mockResolvedValue({
|
||||
MyEndpoint: {
|
||||
type: EModelEndpoint.custom,
|
||||
},
|
||||
});
|
||||
|
||||
const req = createReq(
|
||||
{
|
||||
endpoint: 'MyEndpoint',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
model: 'gpt-4o',
|
||||
temperature: 0.7,
|
||||
region: 'us-east-1',
|
||||
},
|
||||
{ modelSpecs: null },
|
||||
);
|
||||
|
||||
await buildEndpointOption(req, createRes(), jest.fn());
|
||||
|
||||
const parsedResult = parseCompactConvo.mock.results[0].value;
|
||||
expect(parsedResult.region).toBeUndefined();
|
||||
expect(parsedResult.temperature).toBe(0.7);
|
||||
});
|
||||
|
||||
it('should pass defaultParamsEndpoint when re-parsing enforced model spec', async () => {
|
||||
mockGetEndpointsConfig.mockResolvedValue({
|
||||
AnthropicClaude: {
|
||||
type: EModelEndpoint.custom,
|
||||
customParams: {
|
||||
defaultParamsEndpoint: EModelEndpoint.anthropic,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const modelSpec = {
|
||||
name: 'claude-opus-4.5',
|
||||
preset: {
|
||||
endpoint: 'AnthropicClaude',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
model: 'anthropic/claude-opus-4.5',
|
||||
temperature: 0.7,
|
||||
maxOutputTokens: 8192,
|
||||
maxContextTokens: 50000,
|
||||
},
|
||||
};
|
||||
|
||||
const req = createReq(
|
||||
{
|
||||
endpoint: 'AnthropicClaude',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
spec: 'claude-opus-4.5',
|
||||
model: 'anthropic/claude-opus-4.5',
|
||||
},
|
||||
{
|
||||
modelSpecs: {
|
||||
enforce: true,
|
||||
list: [modelSpec],
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
await buildEndpointOption(req, createRes(), jest.fn());
|
||||
|
||||
const enforcedCall = parseCompactConvo.mock.calls[1];
|
||||
expect(enforcedCall[0]).toEqual(
|
||||
expect.objectContaining({
|
||||
defaultParamsEndpoint: EModelEndpoint.anthropic,
|
||||
}),
|
||||
);
|
||||
|
||||
const enforcedResult = parseCompactConvo.mock.results[1].value;
|
||||
expect(enforcedResult.maxOutputTokens).toBe(8192);
|
||||
expect(enforcedResult.temperature).toBe(0.7);
|
||||
expect(enforcedResult.maxContextTokens).toBe(50000);
|
||||
});
|
||||
|
||||
it('should fall back to OpenAI schema when getEndpointsConfig fails', async () => {
|
||||
mockGetEndpointsConfig.mockRejectedValue(new Error('Config unavailable'));
|
||||
|
||||
const req = createReq(
|
||||
{
|
||||
endpoint: 'AnthropicClaude',
|
||||
endpointType: EModelEndpoint.custom,
|
||||
model: 'anthropic/claude-opus-4.5',
|
||||
temperature: 0.7,
|
||||
maxOutputTokens: 8192,
|
||||
max_tokens: 4096,
|
||||
},
|
||||
{ modelSpecs: null },
|
||||
);
|
||||
|
||||
await buildEndpointOption(req, createRes(), jest.fn());
|
||||
|
||||
expect(parseCompactConvo).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
defaultParamsEndpoint: undefined,
|
||||
}),
|
||||
);
|
||||
|
||||
const parsedResult = parseCompactConvo.mock.results[0].value;
|
||||
expect(parsedResult.maxOutputTokens).toBeUndefined();
|
||||
expect(parsedResult.max_tokens).toBe(4096);
|
||||
});
|
||||
});
|
||||
|
|
@ -9,6 +9,7 @@ const resourceToPermissionType = {
|
|||
[ResourceType.AGENT]: PermissionTypes.AGENTS,
|
||||
[ResourceType.PROMPTGROUP]: PermissionTypes.PROMPTS,
|
||||
[ResourceType.MCPSERVER]: PermissionTypes.MCP_SERVERS,
|
||||
[ResourceType.REMOTE_AGENT]: PermissionTypes.REMOTE_AGENTS,
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -7,16 +7,13 @@ const { isEnabled } = require('@librechat/api');
|
|||
* Switches between JWT and OpenID authentication based on cookies and environment settings
|
||||
*/
|
||||
const requireJwtAuth = (req, res, next) => {
|
||||
// Check if token provider is specified in cookies
|
||||
const cookieHeader = req.headers.cookie;
|
||||
const tokenProvider = cookieHeader ? cookies.parse(cookieHeader).token_provider : null;
|
||||
|
||||
// Use OpenID authentication if token provider is OpenID and OPENID_REUSE_TOKENS is enabled
|
||||
if (tokenProvider === 'openid' && isEnabled(process.env.OPENID_REUSE_TOKENS)) {
|
||||
return passport.authenticate('openidJwt', { session: false })(req, res, next);
|
||||
}
|
||||
|
||||
// Default to standard JWT authentication
|
||||
return passport.authenticate('jwt', { session: false })(req, res, next);
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -385,6 +385,40 @@ describe('Convos Routes', () => {
|
|||
expect(deleteConvoSharedLink).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 400 when request body is empty (DoS prevention)', async () => {
|
||||
const response = await request(app).delete('/api/convos').send({});
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body).toEqual({ error: 'no parameters provided' });
|
||||
expect(deleteConvos).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 400 when arg is null (DoS prevention)', async () => {
|
||||
const response = await request(app).delete('/api/convos').send({ arg: null });
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body).toEqual({ error: 'no parameters provided' });
|
||||
expect(deleteConvos).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 400 when arg is undefined (DoS prevention)', async () => {
|
||||
const response = await request(app).delete('/api/convos').send({ arg: undefined });
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body).toEqual({ error: 'no parameters provided' });
|
||||
expect(deleteConvos).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 400 when request body is null (DoS prevention)', async () => {
|
||||
const response = await request(app)
|
||||
.delete('/api/convos')
|
||||
.set('Content-Type', 'application/json')
|
||||
.send('null');
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(deleteConvos).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 500 if deleteConvoSharedLink fails', async () => {
|
||||
const mockConversationId = 'conv-error';
|
||||
|
||||
|
|
|
|||
174
api/server/routes/__tests__/keys.spec.js
Normal file
174
api/server/routes/__tests__/keys.spec.js
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
const express = require('express');
|
||||
const request = require('supertest');
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
updateUserKey: jest.fn(),
|
||||
deleteUserKey: jest.fn(),
|
||||
getUserKeyExpiry: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/middleware/requireJwtAuth', () => (req, res, next) => next());
|
||||
|
||||
jest.mock('~/server/middleware', () => ({
|
||||
requireJwtAuth: (req, res, next) => next(),
|
||||
}));
|
||||
|
||||
describe('Keys Routes', () => {
|
||||
let app;
|
||||
const { updateUserKey, deleteUserKey, getUserKeyExpiry } = require('~/models');
|
||||
|
||||
beforeAll(() => {
|
||||
const keysRouter = require('../keys');
|
||||
|
||||
app = express();
|
||||
app.use(express.json());
|
||||
|
||||
app.use((req, res, next) => {
|
||||
req.user = { id: 'test-user-123' };
|
||||
next();
|
||||
});
|
||||
|
||||
app.use('/api/keys', keysRouter);
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('PUT /', () => {
|
||||
it('should update a user key with the authenticated user ID', async () => {
|
||||
updateUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app)
|
||||
.put('/api/keys')
|
||||
.send({ name: 'openAI', value: 'sk-test-key-123', expiresAt: '2026-12-31' });
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
expect(updateUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'openAI',
|
||||
value: 'sk-test-key-123',
|
||||
expiresAt: '2026-12-31',
|
||||
});
|
||||
expect(updateUserKey).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should not allow userId override via request body (IDOR prevention)', async () => {
|
||||
updateUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app).put('/api/keys').send({
|
||||
userId: 'attacker-injected-id',
|
||||
name: 'openAI',
|
||||
value: 'sk-attacker-key',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
expect(updateUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'openAI',
|
||||
value: 'sk-attacker-key',
|
||||
expiresAt: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('should ignore extraneous fields from request body', async () => {
|
||||
updateUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app).put('/api/keys').send({
|
||||
name: 'openAI',
|
||||
value: 'sk-test-key',
|
||||
expiresAt: '2026-12-31',
|
||||
_id: 'injected-mongo-id',
|
||||
__v: 99,
|
||||
extra: 'should-be-ignored',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
expect(updateUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'openAI',
|
||||
value: 'sk-test-key',
|
||||
expiresAt: '2026-12-31',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle missing optional fields', async () => {
|
||||
updateUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app)
|
||||
.put('/api/keys')
|
||||
.send({ name: 'anthropic', value: 'sk-ant-key' });
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
expect(updateUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'anthropic',
|
||||
value: 'sk-ant-key',
|
||||
expiresAt: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return 400 when request body is null', async () => {
|
||||
const response = await request(app)
|
||||
.put('/api/keys')
|
||||
.set('Content-Type', 'application/json')
|
||||
.send('null');
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(updateUserKey).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('DELETE /:name', () => {
|
||||
it('should delete a user key by name', async () => {
|
||||
deleteUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app).delete('/api/keys/openAI');
|
||||
|
||||
expect(response.status).toBe(204);
|
||||
expect(deleteUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'openAI',
|
||||
});
|
||||
expect(deleteUserKey).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('DELETE /', () => {
|
||||
it('should delete all keys when all=true', async () => {
|
||||
deleteUserKey.mockResolvedValue({});
|
||||
|
||||
const response = await request(app).delete('/api/keys?all=true');
|
||||
|
||||
expect(response.status).toBe(204);
|
||||
expect(deleteUserKey).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
all: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return 400 when all query param is not true', async () => {
|
||||
const response = await request(app).delete('/api/keys');
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body).toEqual({ error: 'Specify either all=true to delete.' });
|
||||
expect(deleteUserKey).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('GET /', () => {
|
||||
it('should return key expiry for a given key name', async () => {
|
||||
const mockExpiry = { expiresAt: '2026-12-31' };
|
||||
getUserKeyExpiry.mockResolvedValue(mockExpiry);
|
||||
|
||||
const response = await request(app).get('/api/keys?name=openAI');
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual(mockExpiry);
|
||||
expect(getUserKeyExpiry).toHaveBeenCalledWith({
|
||||
userId: 'test-user-123',
|
||||
name: 'openAI',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,8 +1,18 @@
|
|||
const crypto = require('crypto');
|
||||
const express = require('express');
|
||||
const request = require('supertest');
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const cookieParser = require('cookie-parser');
|
||||
const { getBasePath } = require('@librechat/api');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
|
||||
function generateTestCsrfToken(flowId) {
|
||||
return crypto
|
||||
.createHmac('sha256', process.env.JWT_SECRET)
|
||||
.update(flowId)
|
||||
.digest('hex')
|
||||
.slice(0, 32);
|
||||
}
|
||||
|
||||
const mockRegistryInstance = {
|
||||
getServerConfig: jest.fn(),
|
||||
|
|
@ -130,6 +140,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
app = express();
|
||||
app.use(express.json());
|
||||
app.use(cookieParser());
|
||||
|
||||
app.use((req, res, next) => {
|
||||
req.user = { id: 'test-user-id' };
|
||||
|
|
@ -168,12 +179,12 @@ describe('MCP Routes', () => {
|
|||
|
||||
MCPOAuthHandler.initiateOAuthFlow.mockResolvedValue({
|
||||
authorizationUrl: 'https://oauth.example.com/auth',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/initiate').query({
|
||||
userId: 'test-user-id',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
|
|
@ -190,7 +201,7 @@ describe('MCP Routes', () => {
|
|||
it('should return 403 when userId does not match authenticated user', async () => {
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/initiate').query({
|
||||
userId: 'different-user-id',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
|
|
@ -228,7 +239,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/initiate').query({
|
||||
userId: 'test-user-id',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
|
|
@ -245,7 +256,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/initiate').query({
|
||||
userId: 'test-user-id',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(500);
|
||||
|
|
@ -255,7 +266,7 @@ describe('MCP Routes', () => {
|
|||
it('should return 400 when flow state metadata is null', async () => {
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({
|
||||
id: 'test-flow-id',
|
||||
id: 'test-user-id:test-server',
|
||||
metadata: null,
|
||||
}),
|
||||
};
|
||||
|
|
@ -265,7 +276,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/initiate').query({
|
||||
userId: 'test-user-id',
|
||||
flowId: 'test-flow-id',
|
||||
flowId: 'test-user-id:test-server',
|
||||
});
|
||||
|
||||
expect(response.status).toBe(400);
|
||||
|
|
@ -280,7 +291,7 @@ describe('MCP Routes', () => {
|
|||
it('should redirect to error page when OAuth error is received', async () => {
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
error: 'access_denied',
|
||||
state: 'test-flow-id',
|
||||
state: 'test-user-id:test-server',
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
|
|
@ -290,7 +301,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
it('should redirect to error page when code is missing', async () => {
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
state: 'test-flow-id',
|
||||
state: 'test-user-id:test-server',
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
|
|
@ -308,15 +319,50 @@ describe('MCP Routes', () => {
|
|||
expect(response.headers.location).toBe(`${basePath}/oauth/error?error=missing_state`);
|
||||
});
|
||||
|
||||
it('should redirect to error page when flow state is not found', async () => {
|
||||
MCPOAuthHandler.getFlowState.mockResolvedValue(null);
|
||||
|
||||
it('should redirect to error page when CSRF cookie is missing', async () => {
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'invalid-flow-id',
|
||||
state: 'test-user-id:test-server',
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(
|
||||
`${basePath}/oauth/error?error=csrf_validation_failed`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should redirect to error page when CSRF cookie does not match state', async () => {
|
||||
const csrfToken = generateTestCsrfToken('different-flow-id');
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-user-id:test-server',
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(
|
||||
`${basePath}/oauth/error?error=csrf_validation_failed`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should redirect to error page when flow state is not found', async () => {
|
||||
MCPOAuthHandler.getFlowState.mockResolvedValue(null);
|
||||
const flowId = 'invalid-flow:id';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/error?error=invalid_state`);
|
||||
});
|
||||
|
|
@ -369,16 +415,22 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
setCachedTools.mockResolvedValue();
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/success?serverName=test-server`);
|
||||
expect(MCPOAuthHandler.completeOAuthFlow).toHaveBeenCalledWith(
|
||||
'test-flow-id',
|
||||
flowId,
|
||||
'test-auth-code',
|
||||
mockFlowManager,
|
||||
{},
|
||||
|
|
@ -400,16 +452,24 @@ describe('MCP Routes', () => {
|
|||
'mcp_oauth',
|
||||
mockTokens,
|
||||
);
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith('test-flow-id', 'mcp_get_tokens');
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith(
|
||||
'test-user-id:test-server',
|
||||
'mcp_get_tokens',
|
||||
);
|
||||
});
|
||||
|
||||
it('should redirect to error page when callback processing fails', async () => {
|
||||
MCPOAuthHandler.getFlowState.mockRejectedValue(new Error('Callback error'));
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
|
|
@ -442,15 +502,21 @@ describe('MCP Routes', () => {
|
|||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/success?serverName=test-server`);
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith('test-flow-id', 'mcp_get_tokens');
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith(flowId, 'mcp_get_tokens');
|
||||
});
|
||||
|
||||
it('should handle reconnection failure after OAuth', async () => {
|
||||
|
|
@ -488,16 +554,22 @@ describe('MCP Routes', () => {
|
|||
getCachedTools.mockResolvedValue({});
|
||||
setCachedTools.mockResolvedValue();
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/success?serverName=test-server`);
|
||||
expect(MCPTokenStorage.storeTokens).toHaveBeenCalled();
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith('test-flow-id', 'mcp_get_tokens');
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith(flowId, 'mcp_get_tokens');
|
||||
});
|
||||
|
||||
it('should redirect to error page if token storage fails', async () => {
|
||||
|
|
@ -530,10 +602,16 @@ describe('MCP Routes', () => {
|
|||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
|
|
@ -589,22 +667,27 @@ describe('MCP Routes', () => {
|
|||
clearReconnection: jest.fn(),
|
||||
});
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/success?serverName=test-server`);
|
||||
|
||||
// Verify storeTokens was called with ORIGINAL flow state credentials
|
||||
expect(MCPTokenStorage.storeTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
userId: 'test-user-id',
|
||||
serverName: 'test-server',
|
||||
tokens: mockTokens,
|
||||
clientInfo: clientInfo, // Uses original flow state, not any "updated" credentials
|
||||
clientInfo: clientInfo,
|
||||
metadata: flowState.metadata,
|
||||
}),
|
||||
);
|
||||
|
|
@ -631,16 +714,21 @@ describe('MCP Routes', () => {
|
|||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
code: 'test-auth-code',
|
||||
state: 'test-flow-id',
|
||||
});
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/success?serverName=test-server`);
|
||||
|
||||
// Verify completeOAuthFlow was NOT called (prevented duplicate)
|
||||
expect(MCPOAuthHandler.completeOAuthFlow).not.toHaveBeenCalled();
|
||||
expect(MCPTokenStorage.storeTokens).not.toHaveBeenCalled();
|
||||
});
|
||||
|
|
@ -755,7 +843,7 @@ describe('MCP Routes', () => {
|
|||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/oauth/status/test-flow-id');
|
||||
const response = await request(app).get('/api/mcp/oauth/status/test-user-id:test-server');
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({
|
||||
|
|
@ -766,6 +854,13 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
});
|
||||
|
||||
it('should return 403 when flowId does not match authenticated user', async () => {
|
||||
const response = await request(app).get('/api/mcp/oauth/status/other-user-id:test-server');
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body).toEqual({ error: 'Access denied' });
|
||||
});
|
||||
|
||||
it('should return 404 when flow is not found', async () => {
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue(null),
|
||||
|
|
@ -774,7 +869,7 @@ describe('MCP Routes', () => {
|
|||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/oauth/status/non-existent-flow');
|
||||
const response = await request(app).get('/api/mcp/oauth/status/test-user-id:non-existent');
|
||||
|
||||
expect(response.status).toBe(404);
|
||||
expect(response.body).toEqual({ error: 'Flow not found' });
|
||||
|
|
@ -788,7 +883,7 @@ describe('MCP Routes', () => {
|
|||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const response = await request(app).get('/api/mcp/oauth/status/error-flow-id');
|
||||
const response = await request(app).get('/api/mcp/oauth/status/test-user-id:error-server');
|
||||
|
||||
expect(response.status).toBe(500);
|
||||
expect(response.body).toEqual({ error: 'Failed to get flow status' });
|
||||
|
|
@ -1375,7 +1470,7 @@ describe('MCP Routes', () => {
|
|||
refresh_token: 'edge-refresh-token',
|
||||
};
|
||||
MCPOAuthHandler.getFlowState = jest.fn().mockResolvedValue({
|
||||
id: 'test-flow-id',
|
||||
id: 'test-user-id:test-server',
|
||||
userId: 'test-user-id',
|
||||
metadata: {
|
||||
serverUrl: 'https://example.com',
|
||||
|
|
@ -1403,8 +1498,12 @@ describe('MCP Routes', () => {
|
|||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback?code=test-code&state=test-flow-id')
|
||||
.get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`)
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.expect(302);
|
||||
|
||||
const basePath = getBasePath();
|
||||
|
|
@ -1424,7 +1523,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({
|
||||
id: 'test-flow-id',
|
||||
id: 'test-user-id:test-server',
|
||||
userId: 'test-user-id',
|
||||
metadata: { serverUrl: 'https://example.com', oauth: {} },
|
||||
clientInfo: {},
|
||||
|
|
@ -1453,8 +1552,12 @@ describe('MCP Routes', () => {
|
|||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback?code=test-code&state=test-flow-id')
|
||||
.get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`)
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.expect(302);
|
||||
|
||||
const basePath = getBasePath();
|
||||
|
|
|
|||
|
|
@ -53,6 +53,12 @@ const checkResourcePermissionAccess = (requiredPermission) => (req, res, next) =
|
|||
requiredPermission,
|
||||
resourceIdParam: 'resourceId',
|
||||
});
|
||||
} else if (resourceType === ResourceType.REMOTE_AGENT) {
|
||||
middleware = canAccessResource({
|
||||
resourceType: ResourceType.REMOTE_AGENT,
|
||||
requiredPermission,
|
||||
resourceIdParam: 'resourceId',
|
||||
});
|
||||
} else if (resourceType === ResourceType.PROMPTGROUP) {
|
||||
middleware = canAccessResource({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
|
|
|
|||
|
|
@ -1,14 +1,47 @@
|
|||
const express = require('express');
|
||||
const jwt = require('jsonwebtoken');
|
||||
const { getAccessToken, getBasePath } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const {
|
||||
getBasePath,
|
||||
getAccessToken,
|
||||
setOAuthSession,
|
||||
validateOAuthCsrf,
|
||||
OAUTH_CSRF_COOKIE,
|
||||
setOAuthCsrfCookie,
|
||||
validateOAuthSession,
|
||||
OAUTH_SESSION_COOKIE,
|
||||
} = require('@librechat/api');
|
||||
const { findToken, updateToken, createToken } = require('~/models');
|
||||
const { requireJwtAuth } = require('~/server/middleware');
|
||||
const { getFlowStateManager } = require('~/config');
|
||||
const { getLogStores } = require('~/cache');
|
||||
|
||||
const router = express.Router();
|
||||
const JWT_SECRET = process.env.JWT_SECRET;
|
||||
const OAUTH_CSRF_COOKIE_PATH = '/api/actions';
|
||||
|
||||
/**
|
||||
* Sets a CSRF cookie binding the action OAuth flow to the current browser session.
|
||||
* Must be called before the user opens the IdP authorization URL.
|
||||
*
|
||||
* @route POST /actions/:action_id/oauth/bind
|
||||
*/
|
||||
router.post('/:action_id/oauth/bind', requireJwtAuth, setOAuthSession, async (req, res) => {
|
||||
try {
|
||||
const { action_id } = req.params;
|
||||
const user = req.user;
|
||||
if (!user?.id) {
|
||||
return res.status(401).json({ error: 'User not authenticated' });
|
||||
}
|
||||
const flowId = `${user.id}:${action_id}`;
|
||||
setOAuthCsrfCookie(res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
res.json({ success: true });
|
||||
} catch (error) {
|
||||
logger.error('[Action OAuth] Failed to set CSRF binding cookie', error);
|
||||
res.status(500).json({ error: 'Failed to bind OAuth flow' });
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* Handles the OAuth callback and exchanges the authorization code for tokens.
|
||||
|
|
@ -45,7 +78,22 @@ router.get('/:action_id/oauth/callback', async (req, res) => {
|
|||
await flowManager.failFlow(identifier, 'oauth', 'Invalid user ID in state parameter');
|
||||
return res.redirect(`${basePath}/oauth/error?error=invalid_state`);
|
||||
}
|
||||
|
||||
identifier = `${decodedState.user}:${action_id}`;
|
||||
|
||||
if (
|
||||
!validateOAuthCsrf(req, res, identifier, OAUTH_CSRF_COOKIE_PATH) &&
|
||||
!validateOAuthSession(req, decodedState.user)
|
||||
) {
|
||||
logger.error('[Action OAuth] CSRF validation failed: no valid CSRF or session cookie', {
|
||||
identifier,
|
||||
hasCsrfCookie: !!req.cookies?.[OAUTH_CSRF_COOKIE],
|
||||
hasSessionCookie: !!req.cookies?.[OAUTH_SESSION_COOKIE],
|
||||
});
|
||||
await flowManager.failFlow(identifier, 'oauth', 'CSRF validation failed');
|
||||
return res.redirect(`${basePath}/oauth/error?error=csrf_validation_failed`);
|
||||
}
|
||||
|
||||
const flowState = await flowManager.getFlowState(identifier, 'oauth');
|
||||
if (!flowState) {
|
||||
throw new Error('OAuth flow not found');
|
||||
|
|
@ -71,7 +119,6 @@ router.get('/:action_id/oauth/callback', async (req, res) => {
|
|||
);
|
||||
await flowManager.completeFlow(identifier, 'oauth', tokenData);
|
||||
|
||||
/** Redirect to React success page */
|
||||
const serverName = flowState.metadata?.action_name || `Action ${action_id}`;
|
||||
const redirectUrl = `${basePath}/oauth/success?serverName=${encodeURIComponent(serverName)}`;
|
||||
res.redirect(redirectUrl);
|
||||
|
|
|
|||
127
api/server/routes/admin/auth.js
Normal file
127
api/server/routes/admin/auth.js
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
const express = require('express');
|
||||
const passport = require('passport');
|
||||
const { randomState } = require('openid-client');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const {
|
||||
requireAdmin,
|
||||
getAdminPanelUrl,
|
||||
exchangeAdminCode,
|
||||
createSetBalanceConfig,
|
||||
} = require('@librechat/api');
|
||||
const { loginController } = require('~/server/controllers/auth/LoginController');
|
||||
const { createOAuthHandler } = require('~/server/controllers/auth/oauth');
|
||||
const { getAppConfig } = require('~/server/services/Config');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
const { getOpenIdConfig } = require('~/strategies');
|
||||
const middleware = require('~/server/middleware');
|
||||
const { Balance } = require('~/db/models');
|
||||
|
||||
const setBalanceConfig = createSetBalanceConfig({
|
||||
getAppConfig,
|
||||
Balance,
|
||||
});
|
||||
|
||||
const router = express.Router();
|
||||
|
||||
router.post(
|
||||
'/login/local',
|
||||
middleware.logHeaders,
|
||||
middleware.loginLimiter,
|
||||
middleware.checkBan,
|
||||
middleware.requireLocalAuth,
|
||||
requireAdmin,
|
||||
setBalanceConfig,
|
||||
loginController,
|
||||
);
|
||||
|
||||
router.get('/verify', middleware.requireJwtAuth, requireAdmin, (req, res) => {
|
||||
const { password: _p, totpSecret: _t, __v, ...user } = req.user;
|
||||
user.id = user._id.toString();
|
||||
res.status(200).json({ user });
|
||||
});
|
||||
|
||||
router.get('/oauth/openid/check', (req, res) => {
|
||||
const openidConfig = getOpenIdConfig();
|
||||
if (!openidConfig) {
|
||||
return res.status(404).json({
|
||||
error: 'OpenID configuration not found',
|
||||
error_code: 'OPENID_NOT_CONFIGURED',
|
||||
});
|
||||
}
|
||||
res.status(200).json({ message: 'OpenID check successful' });
|
||||
});
|
||||
|
||||
router.get('/oauth/openid', (req, res, next) => {
|
||||
return passport.authenticate('openidAdmin', {
|
||||
session: false,
|
||||
state: randomState(),
|
||||
})(req, res, next);
|
||||
});
|
||||
|
||||
router.get(
|
||||
'/oauth/openid/callback',
|
||||
passport.authenticate('openidAdmin', {
|
||||
failureRedirect: `${getAdminPanelUrl()}/auth/openid/callback?error=auth_failed&error_description=Authentication+failed`,
|
||||
failureMessage: true,
|
||||
session: false,
|
||||
}),
|
||||
requireAdmin,
|
||||
setBalanceConfig,
|
||||
middleware.checkDomainAllowed,
|
||||
createOAuthHandler(`${getAdminPanelUrl()}/auth/openid/callback`),
|
||||
);
|
||||
|
||||
/** Regex pattern for valid exchange codes: 64 hex characters */
|
||||
const EXCHANGE_CODE_PATTERN = /^[a-f0-9]{64}$/i;
|
||||
|
||||
/**
|
||||
* Exchange OAuth authorization code for tokens.
|
||||
* This endpoint is called server-to-server by the admin panel.
|
||||
* The code is one-time-use and expires in 30 seconds.
|
||||
*
|
||||
* POST /api/admin/oauth/exchange
|
||||
* Body: { code: string }
|
||||
* Response: { token: string, refreshToken: string, user: object }
|
||||
*/
|
||||
router.post('/oauth/exchange', middleware.loginLimiter, async (req, res) => {
|
||||
try {
|
||||
const { code } = req.body;
|
||||
|
||||
if (!code) {
|
||||
logger.warn('[admin/oauth/exchange] Missing authorization code');
|
||||
return res.status(400).json({
|
||||
error: 'Missing authorization code',
|
||||
error_code: 'MISSING_CODE',
|
||||
});
|
||||
}
|
||||
|
||||
if (typeof code !== 'string' || !EXCHANGE_CODE_PATTERN.test(code)) {
|
||||
logger.warn('[admin/oauth/exchange] Invalid authorization code format');
|
||||
return res.status(400).json({
|
||||
error: 'Invalid authorization code format',
|
||||
error_code: 'INVALID_CODE_FORMAT',
|
||||
});
|
||||
}
|
||||
|
||||
const cache = getLogStores(CacheKeys.ADMIN_OAUTH_EXCHANGE);
|
||||
const result = await exchangeAdminCode(cache, code);
|
||||
|
||||
if (!result) {
|
||||
return res.status(401).json({
|
||||
error: 'Invalid or expired authorization code',
|
||||
error_code: 'INVALID_OR_EXPIRED_CODE',
|
||||
});
|
||||
}
|
||||
|
||||
res.json(result);
|
||||
} catch (error) {
|
||||
logger.error('[admin/oauth/exchange] Error:', error);
|
||||
res.status(500).json({
|
||||
error: 'Internal server error',
|
||||
error_code: 'INTERNAL_ERROR',
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
module.exports = router;
|
||||
303
api/server/routes/agents/__tests__/abort.spec.js
Normal file
303
api/server/routes/agents/__tests__/abort.spec.js
Normal file
|
|
@ -0,0 +1,303 @@
|
|||
/**
|
||||
* Tests for the agent abort endpoint
|
||||
*
|
||||
* Tests the following fixes from PR #11462:
|
||||
* 1. Authorization check - only job owner can abort
|
||||
* 2. Early abort handling - skip save when no responseMessageId
|
||||
* 3. Partial response saving - save message before returning
|
||||
*/
|
||||
|
||||
const express = require('express');
|
||||
const request = require('supertest');
|
||||
|
||||
const mockLogger = {
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
};
|
||||
|
||||
const mockGenerationJobManager = {
|
||||
getJob: jest.fn(),
|
||||
abortJob: jest.fn(),
|
||||
getActiveJobIdsForUser: jest.fn(),
|
||||
};
|
||||
|
||||
const mockSaveMessage = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
...jest.requireActual('@librechat/data-schemas'),
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
...jest.requireActual('@librechat/api'),
|
||||
isEnabled: jest.fn().mockReturnValue(false),
|
||||
GenerationJobManager: mockGenerationJobManager,
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: (...args) => mockSaveMessage(...args),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/middleware', () => ({
|
||||
uaParser: (req, res, next) => next(),
|
||||
checkBan: (req, res, next) => next(),
|
||||
requireJwtAuth: (req, res, next) => {
|
||||
req.user = { id: 'test-user-123' };
|
||||
next();
|
||||
},
|
||||
messageIpLimiter: (req, res, next) => next(),
|
||||
configMiddleware: (req, res, next) => next(),
|
||||
messageUserLimiter: (req, res, next) => next(),
|
||||
}));
|
||||
|
||||
// Mock the chat module - needs to be a router
|
||||
jest.mock('~/server/routes/agents/chat', () => require('express').Router());
|
||||
|
||||
// Mock the v1 module - v1 is directly used as middleware
|
||||
jest.mock('~/server/routes/agents/v1', () => ({
|
||||
v1: require('express').Router(),
|
||||
}));
|
||||
|
||||
// Import after mocks
|
||||
const agentRoutes = require('~/server/routes/agents/index');
|
||||
|
||||
describe('Agent Abort Endpoint', () => {
|
||||
let app;
|
||||
|
||||
beforeAll(() => {
|
||||
app = express();
|
||||
app.use(express.json());
|
||||
app.use('/api/agents', agentRoutes);
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('POST /chat/abort', () => {
|
||||
describe('Authorization', () => {
|
||||
it("should return 403 when user tries to abort another user's job", async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'other-user-456' },
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body).toEqual({ error: 'Unauthorized' });
|
||||
expect(mockLogger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Unauthorized abort attempt'),
|
||||
);
|
||||
expect(mockGenerationJobManager.abortJob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should allow abort when user owns the job', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: null,
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith(jobStreamId);
|
||||
});
|
||||
|
||||
it('should allow abort when job has no userId metadata (backwards compatibility)', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: {},
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: null,
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
});
|
||||
});
|
||||
|
||||
describe('Early Abort Handling', () => {
|
||||
it('should skip message saving when responseMessageId is missing (early abort)', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: 'user-msg-123' },
|
||||
// No responseMessageId - early abort before generation started
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip message saving when userMessage is missing', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
// No userMessage
|
||||
responseMessageId: 'response-msg-123',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Partial Response Saving', () => {
|
||||
it('should save partial response when both userMessage and responseMessageId exist', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
const userMessageId = 'user-msg-123';
|
||||
const responseMessageId = 'response-msg-456';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: userMessageId },
|
||||
responseMessageId,
|
||||
conversationId: jobStreamId,
|
||||
sender: 'TestAgent',
|
||||
endpoint: 'anthropic',
|
||||
model: 'claude-3',
|
||||
},
|
||||
content: [{ type: 'text', text: 'Partial response...' }],
|
||||
text: 'Partial response...',
|
||||
});
|
||||
|
||||
mockSaveMessage.mockResolvedValue();
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
expect.objectContaining({
|
||||
messageId: responseMessageId,
|
||||
parentMessageId: userMessageId,
|
||||
conversationId: jobStreamId,
|
||||
content: [{ type: 'text', text: 'Partial response...' }],
|
||||
text: 'Partial response...',
|
||||
sender: 'TestAgent',
|
||||
endpoint: 'anthropic',
|
||||
model: 'claude-3',
|
||||
unfinished: true,
|
||||
error: false,
|
||||
isCreatedByUser: false,
|
||||
user: 'test-user-123',
|
||||
}),
|
||||
expect.objectContaining({
|
||||
context: 'api/server/routes/agents/index.js - abort endpoint',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle saveMessage errors gracefully', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: 'user-msg-123' },
|
||||
responseMessageId: 'response-msg-456',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
mockSaveMessage.mockRejectedValue(new Error('Database error'));
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
// Should still return success even if save fails
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
expect(mockLogger.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Failed to save partial response'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Job Not Found', () => {
|
||||
it('should return 404 when job is not found', async () => {
|
||||
mockGenerationJobManager.getJob.mockResolvedValue(null);
|
||||
mockGenerationJobManager.getActiveJobIdsForUser.mockResolvedValue([]);
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: 'non-existent-job' });
|
||||
|
||||
expect(response.status).toBe(404);
|
||||
expect(response.body).toEqual({
|
||||
error: 'Job not found',
|
||||
streamId: 'non-existent-job',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
1125
api/server/routes/agents/__tests__/responses.spec.js
Normal file
1125
api/server/routes/agents/__tests__/responses.spec.js
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -9,6 +9,9 @@ const {
|
|||
configMiddleware,
|
||||
messageUserLimiter,
|
||||
} = require('~/server/middleware');
|
||||
const { saveMessage } = require('~/models');
|
||||
const openai = require('./openai');
|
||||
const responses = require('./responses');
|
||||
const { v1 } = require('./v1');
|
||||
const chat = require('./chat');
|
||||
|
||||
|
|
@ -16,6 +19,20 @@ const { LIMIT_MESSAGE_IP, LIMIT_MESSAGE_USER } = process.env ?? {};
|
|||
|
||||
const router = express.Router();
|
||||
|
||||
/**
|
||||
* Open Responses API routes (API key authentication handled in route file)
|
||||
* Mounted at /agents/v1/responses (full path: /api/agents/v1/responses)
|
||||
* NOTE: Must be mounted BEFORE /v1 to avoid being caught by the less specific route
|
||||
* @see https://openresponses.org/specification
|
||||
*/
|
||||
router.use('/v1/responses', responses);
|
||||
|
||||
/**
|
||||
* OpenAI-compatible API routes (API key authentication handled in route file)
|
||||
* Mounted at /agents/v1 (full path: /api/agents/v1/chat/completions)
|
||||
*/
|
||||
router.use('/v1', openai);
|
||||
|
||||
router.use(requireJwtAuth);
|
||||
router.use(checkBan);
|
||||
router.use(uaParser);
|
||||
|
|
@ -46,6 +63,10 @@ router.get('/chat/stream/:streamId', async (req, res) => {
|
|||
});
|
||||
}
|
||||
|
||||
if (job.metadata?.userId && job.metadata.userId !== req.user.id) {
|
||||
return res.status(403).json({ error: 'Unauthorized' });
|
||||
}
|
||||
|
||||
res.setHeader('Content-Encoding', 'identity');
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-cache, no-transform');
|
||||
|
|
@ -194,9 +215,53 @@ router.post('/chat/abort', async (req, res) => {
|
|||
logger.debug(`[AgentStream] Computed jobStreamId: ${jobStreamId}`);
|
||||
|
||||
if (job && jobStreamId) {
|
||||
if (job.metadata?.userId && job.metadata.userId !== userId) {
|
||||
logger.warn(`[AgentStream] Unauthorized abort attempt for ${jobStreamId} by user ${userId}`);
|
||||
return res.status(403).json({ error: 'Unauthorized' });
|
||||
}
|
||||
|
||||
logger.debug(`[AgentStream] Job found, aborting: ${jobStreamId}`);
|
||||
await GenerationJobManager.abortJob(jobStreamId);
|
||||
logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`);
|
||||
const abortResult = await GenerationJobManager.abortJob(jobStreamId);
|
||||
logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`, {
|
||||
abortResultSuccess: abortResult.success,
|
||||
abortResultUserMessageId: abortResult.jobData?.userMessage?.messageId,
|
||||
abortResultResponseMessageId: abortResult.jobData?.responseMessageId,
|
||||
});
|
||||
|
||||
// CRITICAL: Save partial response BEFORE returning to prevent race condition.
|
||||
// If user sends a follow-up immediately after abort, the parentMessageId must exist in DB.
|
||||
// Only save if we have a valid responseMessageId (skip early aborts before generation started)
|
||||
if (
|
||||
abortResult.success &&
|
||||
abortResult.jobData?.userMessage?.messageId &&
|
||||
abortResult.jobData?.responseMessageId
|
||||
) {
|
||||
const { jobData, content, text } = abortResult;
|
||||
const responseMessage = {
|
||||
messageId: jobData.responseMessageId,
|
||||
parentMessageId: jobData.userMessage.messageId,
|
||||
conversationId: jobData.conversationId,
|
||||
content: content || [],
|
||||
text: text || '',
|
||||
sender: jobData.sender || 'AI',
|
||||
endpoint: jobData.endpoint,
|
||||
model: jobData.model,
|
||||
unfinished: true,
|
||||
error: false,
|
||||
isCreatedByUser: false,
|
||||
user: userId,
|
||||
};
|
||||
|
||||
try {
|
||||
await saveMessage(req, responseMessage, {
|
||||
context: 'api/server/routes/agents/index.js - abort endpoint',
|
||||
});
|
||||
logger.debug(`[AgentStream] Saved partial response for: ${jobStreamId}`);
|
||||
} catch (saveError) {
|
||||
logger.error(`[AgentStream] Failed to save partial response: ${saveError.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
return res.json({ success: true, aborted: jobStreamId });
|
||||
}
|
||||
|
||||
|
|
|
|||
110
api/server/routes/agents/openai.js
Normal file
110
api/server/routes/agents/openai.js
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
/**
|
||||
* OpenAI-compatible API routes for LibreChat agents.
|
||||
*
|
||||
* Provides a /v1/chat/completions compatible interface for
|
||||
* interacting with LibreChat agents remotely via API.
|
||||
*
|
||||
* Usage:
|
||||
* POST /v1/chat/completions - Chat with an agent
|
||||
* GET /v1/models - List available agents
|
||||
* GET /v1/models/:model - Get agent details
|
||||
*
|
||||
* Request format:
|
||||
* {
|
||||
* "model": "agent_id_here",
|
||||
* "messages": [{"role": "user", "content": "Hello!"}],
|
||||
* "stream": true
|
||||
* }
|
||||
*/
|
||||
const express = require('express');
|
||||
const { PermissionTypes, Permissions } = require('librechat-data-provider');
|
||||
const {
|
||||
generateCheckAccess,
|
||||
createRequireApiKeyAuth,
|
||||
createCheckRemoteAgentAccess,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
OpenAIChatCompletionController,
|
||||
ListModelsController,
|
||||
GetModelController,
|
||||
} = require('~/server/controllers/agents/openai');
|
||||
const { getEffectivePermissions } = require('~/server/services/PermissionService');
|
||||
const { validateAgentApiKey, findUser } = require('~/models');
|
||||
const { configMiddleware } = require('~/server/middleware');
|
||||
const { getRoleByName } = require('~/models/Role');
|
||||
const { getAgent } = require('~/models/Agent');
|
||||
|
||||
const router = express.Router();
|
||||
|
||||
const requireApiKeyAuth = createRequireApiKeyAuth({
|
||||
validateAgentApiKey,
|
||||
findUser,
|
||||
});
|
||||
|
||||
const checkRemoteAgentsFeature = generateCheckAccess({
|
||||
permissionType: PermissionTypes.REMOTE_AGENTS,
|
||||
permissions: [Permissions.USE],
|
||||
getRoleByName,
|
||||
});
|
||||
|
||||
const checkAgentPermission = createCheckRemoteAgentAccess({
|
||||
getAgent,
|
||||
getEffectivePermissions,
|
||||
});
|
||||
|
||||
router.use(requireApiKeyAuth);
|
||||
router.use(configMiddleware);
|
||||
router.use(checkRemoteAgentsFeature);
|
||||
|
||||
/**
|
||||
* @route POST /v1/chat/completions
|
||||
* @desc OpenAI-compatible chat completions with agents
|
||||
* @access Private (API key auth required)
|
||||
*
|
||||
* Request body:
|
||||
* {
|
||||
* "model": "agent_id", // Required: The agent ID to use
|
||||
* "messages": [...], // Required: Array of chat messages
|
||||
* "stream": true, // Optional: Whether to stream (default: false)
|
||||
* "conversation_id": "...", // Optional: Conversation ID for context
|
||||
* "parent_message_id": "..." // Optional: Parent message for threading
|
||||
* }
|
||||
*
|
||||
* Response (streaming):
|
||||
* - SSE stream with OpenAI chat.completion.chunk format
|
||||
* - Includes delta.reasoning for thinking/reasoning content
|
||||
*
|
||||
* Response (non-streaming):
|
||||
* - Standard OpenAI chat.completion format
|
||||
*/
|
||||
router.post('/chat/completions', checkAgentPermission, OpenAIChatCompletionController);
|
||||
|
||||
/**
|
||||
* @route GET /v1/models
|
||||
* @desc List available agents as models
|
||||
* @access Private (API key auth required)
|
||||
*
|
||||
* Response:
|
||||
* {
|
||||
* "object": "list",
|
||||
* "data": [
|
||||
* {
|
||||
* "id": "agent_id",
|
||||
* "object": "model",
|
||||
* "name": "Agent Name",
|
||||
* "provider": "openai",
|
||||
* ...
|
||||
* }
|
||||
* ]
|
||||
* }
|
||||
*/
|
||||
router.get('/models', ListModelsController);
|
||||
|
||||
/**
|
||||
* @route GET /v1/models/:model
|
||||
* @desc Get details for a specific agent/model
|
||||
* @access Private (API key auth required)
|
||||
*/
|
||||
router.get('/models/:model', GetModelController);
|
||||
|
||||
module.exports = router;
|
||||
144
api/server/routes/agents/responses.js
Normal file
144
api/server/routes/agents/responses.js
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
/**
|
||||
* Open Responses API routes for LibreChat agents.
|
||||
*
|
||||
* Implements the Open Responses specification for a forward-looking,
|
||||
* agentic API that uses items as the fundamental unit and semantic
|
||||
* streaming events.
|
||||
*
|
||||
* Usage:
|
||||
* POST /v1/responses - Create a response
|
||||
* GET /v1/models - List available agents
|
||||
*
|
||||
* Request format:
|
||||
* {
|
||||
* "model": "agent_id_here",
|
||||
* "input": "Hello!" or [{ type: "message", role: "user", content: "Hello!" }],
|
||||
* "stream": true,
|
||||
* "previous_response_id": "optional_conversation_id"
|
||||
* }
|
||||
*
|
||||
* @see https://openresponses.org/specification
|
||||
*/
|
||||
const express = require('express');
|
||||
const { PermissionTypes, Permissions } = require('librechat-data-provider');
|
||||
const {
|
||||
generateCheckAccess,
|
||||
createRequireApiKeyAuth,
|
||||
createCheckRemoteAgentAccess,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
createResponse,
|
||||
getResponse,
|
||||
listModels,
|
||||
} = require('~/server/controllers/agents/responses');
|
||||
const { getEffectivePermissions } = require('~/server/services/PermissionService');
|
||||
const { validateAgentApiKey, findUser } = require('~/models');
|
||||
const { configMiddleware } = require('~/server/middleware');
|
||||
const { getRoleByName } = require('~/models/Role');
|
||||
const { getAgent } = require('~/models/Agent');
|
||||
|
||||
const router = express.Router();
|
||||
|
||||
const requireApiKeyAuth = createRequireApiKeyAuth({
|
||||
validateAgentApiKey,
|
||||
findUser,
|
||||
});
|
||||
|
||||
const checkRemoteAgentsFeature = generateCheckAccess({
|
||||
permissionType: PermissionTypes.REMOTE_AGENTS,
|
||||
permissions: [Permissions.USE],
|
||||
getRoleByName,
|
||||
});
|
||||
|
||||
const checkAgentPermission = createCheckRemoteAgentAccess({
|
||||
getAgent,
|
||||
getEffectivePermissions,
|
||||
});
|
||||
|
||||
router.use(requireApiKeyAuth);
|
||||
router.use(configMiddleware);
|
||||
router.use(checkRemoteAgentsFeature);
|
||||
|
||||
/**
|
||||
* @route POST /v1/responses
|
||||
* @desc Create a model response following Open Responses specification
|
||||
* @access Private (API key auth required)
|
||||
*
|
||||
* Request body:
|
||||
* {
|
||||
* "model": "agent_id", // Required: The agent ID to use
|
||||
* "input": "..." | [...], // Required: String or array of input items
|
||||
* "stream": true, // Optional: Whether to stream (default: false)
|
||||
* "previous_response_id": "...", // Optional: Previous response for continuation
|
||||
* "instructions": "...", // Optional: Additional instructions
|
||||
* "tools": [...], // Optional: Additional tools
|
||||
* "tool_choice": "auto", // Optional: Tool choice mode
|
||||
* "max_output_tokens": 4096, // Optional: Max tokens
|
||||
* "temperature": 0.7 // Optional: Temperature
|
||||
* }
|
||||
*
|
||||
* Response (streaming):
|
||||
* - SSE stream with semantic events:
|
||||
* - response.in_progress
|
||||
* - response.output_item.added
|
||||
* - response.content_part.added
|
||||
* - response.output_text.delta
|
||||
* - response.output_text.done
|
||||
* - response.function_call_arguments.delta
|
||||
* - response.output_item.done
|
||||
* - response.completed
|
||||
* - [DONE]
|
||||
*
|
||||
* Response (non-streaming):
|
||||
* {
|
||||
* "id": "resp_xxx",
|
||||
* "object": "response",
|
||||
* "created_at": 1234567890,
|
||||
* "status": "completed",
|
||||
* "model": "agent_id",
|
||||
* "output": [...], // Array of output items
|
||||
* "usage": { ... }
|
||||
* }
|
||||
*/
|
||||
router.post('/', checkAgentPermission, createResponse);
|
||||
|
||||
/**
|
||||
* @route GET /v1/responses/models
|
||||
* @desc List available agents as models
|
||||
* @access Private (API key auth required)
|
||||
*
|
||||
* Response:
|
||||
* {
|
||||
* "object": "list",
|
||||
* "data": [
|
||||
* {
|
||||
* "id": "agent_id",
|
||||
* "object": "model",
|
||||
* "name": "Agent Name",
|
||||
* "provider": "openai",
|
||||
* ...
|
||||
* }
|
||||
* ]
|
||||
* }
|
||||
*/
|
||||
router.get('/models', listModels);
|
||||
|
||||
/**
|
||||
* @route GET /v1/responses/:id
|
||||
* @desc Retrieve a stored response by ID
|
||||
* @access Private (API key auth required)
|
||||
*
|
||||
* Response:
|
||||
* {
|
||||
* "id": "resp_xxx",
|
||||
* "object": "response",
|
||||
* "created_at": 1234567890,
|
||||
* "status": "completed",
|
||||
* "model": "agent_id",
|
||||
* "output": [...],
|
||||
* "usage": { ... }
|
||||
* }
|
||||
*/
|
||||
router.get('/:id', getResponse);
|
||||
|
||||
module.exports = router;
|
||||
36
api/server/routes/apiKeys.js
Normal file
36
api/server/routes/apiKeys.js
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
const express = require('express');
|
||||
const { generateCheckAccess, createApiKeyHandlers } = require('@librechat/api');
|
||||
const { PermissionTypes, Permissions } = require('librechat-data-provider');
|
||||
const {
|
||||
getAgentApiKeyById,
|
||||
createAgentApiKey,
|
||||
deleteAgentApiKey,
|
||||
listAgentApiKeys,
|
||||
} = require('~/models');
|
||||
const { requireJwtAuth } = require('~/server/middleware');
|
||||
const { getRoleByName } = require('~/models/Role');
|
||||
|
||||
const router = express.Router();
|
||||
|
||||
const handlers = createApiKeyHandlers({
|
||||
createAgentApiKey,
|
||||
listAgentApiKeys,
|
||||
deleteAgentApiKey,
|
||||
getAgentApiKeyById,
|
||||
});
|
||||
|
||||
const checkRemoteAgentsUse = generateCheckAccess({
|
||||
permissionType: PermissionTypes.REMOTE_AGENTS,
|
||||
permissions: [Permissions.USE],
|
||||
getRoleByName,
|
||||
});
|
||||
|
||||
router.post('/', requireJwtAuth, checkRemoteAgentsUse, handlers.createApiKey);
|
||||
|
||||
router.get('/', requireJwtAuth, checkRemoteAgentsUse, handlers.listApiKeys);
|
||||
|
||||
router.get('/:id', requireJwtAuth, checkRemoteAgentsUse, handlers.getApiKey);
|
||||
|
||||
router.delete('/:id', requireJwtAuth, checkRemoteAgentsUse, handlers.deleteApiKey);
|
||||
|
||||
module.exports = router;
|
||||
|
|
@ -98,7 +98,7 @@ router.get('/gen_title/:conversationId', async (req, res) => {
|
|||
|
||||
router.delete('/', async (req, res) => {
|
||||
let filter = {};
|
||||
const { conversationId, source, thread_id, endpoint } = req.body.arg;
|
||||
const { conversationId, source, thread_id, endpoint } = req.body?.arg ?? {};
|
||||
|
||||
// Prevent deletion of all conversations
|
||||
if (!conversationId && !source && !thread_id && !endpoint) {
|
||||
|
|
@ -160,7 +160,7 @@ router.delete('/all', async (req, res) => {
|
|||
* @returns {object} 200 - The updated conversation object.
|
||||
*/
|
||||
router.post('/archive', validateConvoAccess, async (req, res) => {
|
||||
const { conversationId, isArchived } = req.body.arg ?? {};
|
||||
const { conversationId, isArchived } = req.body?.arg ?? {};
|
||||
|
||||
if (!conversationId) {
|
||||
return res.status(400).json({ error: 'conversationId is required' });
|
||||
|
|
@ -194,7 +194,7 @@ const MAX_CONVO_TITLE_LENGTH = 1024;
|
|||
* @returns {object} 201 - The updated conversation object.
|
||||
*/
|
||||
router.post('/update', validateConvoAccess, async (req, res) => {
|
||||
const { conversationId, title } = req.body.arg ?? {};
|
||||
const { conversationId, title } = req.body?.arg ?? {};
|
||||
|
||||
if (!conversationId) {
|
||||
return res.status(400).json({ error: 'conversationId is required' });
|
||||
|
|
|
|||
|
|
@ -2,11 +2,11 @@ const path = require('path');
|
|||
const fs = require('fs').promises;
|
||||
const express = require('express');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { isAgentsEndpoint } = require('librechat-data-provider');
|
||||
const { isAssistantsEndpoint } = require('librechat-data-provider');
|
||||
const {
|
||||
filterFile,
|
||||
processImageFile,
|
||||
processAgentFileUpload,
|
||||
processImageFile,
|
||||
filterFile,
|
||||
} = require('~/server/services/Files/process');
|
||||
|
||||
const router = express.Router();
|
||||
|
|
@ -21,7 +21,7 @@ router.post('/', async (req, res) => {
|
|||
metadata.temp_file_id = metadata.file_id;
|
||||
metadata.file_id = req.file_id;
|
||||
|
||||
if (isAgentsEndpoint(metadata.endpoint) && metadata.tool_resource != null) {
|
||||
if (!isAssistantsEndpoint(metadata.endpoint) && metadata.tool_resource != null) {
|
||||
return await processAgentFileUpload({ req, res, metadata });
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
const accessPermissions = require('./accessPermissions');
|
||||
const assistants = require('./assistants');
|
||||
const categories = require('./categories');
|
||||
const adminAuth = require('./admin/auth');
|
||||
const endpoints = require('./endpoints');
|
||||
const staticRoute = require('./static');
|
||||
const messages = require('./messages');
|
||||
|
|
@ -9,6 +10,7 @@ const presets = require('./presets');
|
|||
const prompts = require('./prompts');
|
||||
const balance = require('./balance');
|
||||
const actions = require('./actions');
|
||||
const apiKeys = require('./apiKeys');
|
||||
const banner = require('./banner');
|
||||
const search = require('./search');
|
||||
const models = require('./models');
|
||||
|
|
@ -28,7 +30,9 @@ const mcp = require('./mcp');
|
|||
module.exports = {
|
||||
mcp,
|
||||
auth,
|
||||
adminAuth,
|
||||
keys,
|
||||
apiKeys,
|
||||
user,
|
||||
tags,
|
||||
roles,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,11 @@ const { requireJwtAuth } = require('~/server/middleware');
|
|||
const router = express.Router();
|
||||
|
||||
router.put('/', requireJwtAuth, async (req, res) => {
|
||||
await updateUserKey({ userId: req.user.id, ...req.body });
|
||||
if (req.body == null || typeof req.body !== 'object') {
|
||||
return res.status(400).send({ error: 'Invalid request body.' });
|
||||
}
|
||||
const { name, value, expiresAt } = req.body;
|
||||
await updateUserKey({ userId: req.user.id, name, value, expiresAt });
|
||||
res.status(201).send();
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -8,18 +8,32 @@ const {
|
|||
Permissions,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
getBasePath,
|
||||
createSafeUser,
|
||||
MCPOAuthHandler,
|
||||
MCPTokenStorage,
|
||||
getBasePath,
|
||||
setOAuthSession,
|
||||
getUserMCPAuthMap,
|
||||
validateOAuthCsrf,
|
||||
OAUTH_CSRF_COOKIE,
|
||||
setOAuthCsrfCookie,
|
||||
generateCheckAccess,
|
||||
validateOAuthSession,
|
||||
OAUTH_SESSION_COOKIE,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
getMCPManager,
|
||||
getFlowStateManager,
|
||||
createMCPServerController,
|
||||
updateMCPServerController,
|
||||
deleteMCPServerController,
|
||||
getMCPServersList,
|
||||
getMCPServerById,
|
||||
getMCPTools,
|
||||
} = require('~/server/controllers/mcp');
|
||||
const {
|
||||
getOAuthReconnectionManager,
|
||||
getMCPServersRegistry,
|
||||
getFlowStateManager,
|
||||
getMCPManager,
|
||||
} = require('~/config');
|
||||
const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
|
||||
const { requireJwtAuth, canAccessMCPServerResource } = require('~/server/middleware');
|
||||
|
|
@ -27,20 +41,14 @@ const { findToken, updateToken, createToken, deleteTokens } = require('~/models'
|
|||
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
|
||||
const { updateMCPServerTools } = require('~/server/services/Config/mcp');
|
||||
const { reinitMCPServer } = require('~/server/services/Tools/mcp');
|
||||
const { getMCPTools } = require('~/server/controllers/mcp');
|
||||
const { findPluginAuthsByKeys } = require('~/models');
|
||||
const { getRoleByName } = require('~/models/Role');
|
||||
const { getLogStores } = require('~/cache');
|
||||
const {
|
||||
createMCPServerController,
|
||||
getMCPServerById,
|
||||
getMCPServersList,
|
||||
updateMCPServerController,
|
||||
deleteMCPServerController,
|
||||
} = require('~/server/controllers/mcp');
|
||||
|
||||
const router = Router();
|
||||
|
||||
const OAUTH_CSRF_COOKIE_PATH = '/api/mcp';
|
||||
|
||||
/**
|
||||
* Get all MCP tools available to the user
|
||||
* Returns only MCP tools, completely decoupled from regular LibreChat tools
|
||||
|
|
@ -53,7 +61,7 @@ router.get('/tools', requireJwtAuth, async (req, res) => {
|
|||
* Initiate OAuth flow
|
||||
* This endpoint is called when the user clicks the auth link in the UI
|
||||
*/
|
||||
router.get('/:serverName/oauth/initiate', requireJwtAuth, async (req, res) => {
|
||||
router.get('/:serverName/oauth/initiate', requireJwtAuth, setOAuthSession, async (req, res) => {
|
||||
try {
|
||||
const { serverName } = req.params;
|
||||
const { userId, flowId } = req.query;
|
||||
|
|
@ -93,7 +101,7 @@ router.get('/:serverName/oauth/initiate', requireJwtAuth, async (req, res) => {
|
|||
|
||||
logger.debug('[MCP OAuth] OAuth flow initiated', { oauthFlowId, authorizationUrl });
|
||||
|
||||
// Redirect user to the authorization URL
|
||||
setOAuthCsrfCookie(res, oauthFlowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
res.redirect(authorizationUrl);
|
||||
} catch (error) {
|
||||
logger.error('[MCP OAuth] Failed to initiate OAuth', error);
|
||||
|
|
@ -138,6 +146,25 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
|||
const flowId = state;
|
||||
logger.debug('[MCP OAuth] Using flow ID from state', { flowId });
|
||||
|
||||
const flowParts = flowId.split(':');
|
||||
if (flowParts.length < 2 || !flowParts[0] || !flowParts[1]) {
|
||||
logger.error('[MCP OAuth] Invalid flow ID format in state', { flowId });
|
||||
return res.redirect(`${basePath}/oauth/error?error=invalid_state`);
|
||||
}
|
||||
|
||||
const [flowUserId] = flowParts;
|
||||
if (
|
||||
!validateOAuthCsrf(req, res, flowId, OAUTH_CSRF_COOKIE_PATH) &&
|
||||
!validateOAuthSession(req, flowUserId)
|
||||
) {
|
||||
logger.error('[MCP OAuth] CSRF validation failed: no valid CSRF or session cookie', {
|
||||
flowId,
|
||||
hasCsrfCookie: !!req.cookies?.[OAUTH_CSRF_COOKIE],
|
||||
hasSessionCookie: !!req.cookies?.[OAUTH_SESSION_COOKIE],
|
||||
});
|
||||
return res.redirect(`${basePath}/oauth/error?error=csrf_validation_failed`);
|
||||
}
|
||||
|
||||
const flowsCache = getLogStores(CacheKeys.FLOWS);
|
||||
const flowManager = getFlowStateManager(flowsCache);
|
||||
|
||||
|
|
@ -302,13 +329,47 @@ router.get('/oauth/tokens/:flowId', requireJwtAuth, async (req, res) => {
|
|||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* Set CSRF binding cookie for OAuth flows initiated outside of HTTP request/response
|
||||
* (e.g. during chat via SSE). The frontend should call this before opening the OAuth URL
|
||||
* so the callback can verify the browser matches the flow initiator.
|
||||
*/
|
||||
router.post('/:serverName/oauth/bind', requireJwtAuth, setOAuthSession, async (req, res) => {
|
||||
try {
|
||||
const { serverName } = req.params;
|
||||
const user = req.user;
|
||||
|
||||
if (!user?.id) {
|
||||
return res.status(401).json({ error: 'User not authenticated' });
|
||||
}
|
||||
|
||||
const flowId = MCPOAuthHandler.generateFlowId(user.id, serverName);
|
||||
setOAuthCsrfCookie(res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
|
||||
res.json({ success: true });
|
||||
} catch (error) {
|
||||
logger.error('[MCP OAuth] Failed to set CSRF binding cookie', error);
|
||||
res.status(500).json({ error: 'Failed to bind OAuth flow' });
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* Check OAuth flow status
|
||||
* This endpoint can be used to poll the status of an OAuth flow
|
||||
*/
|
||||
router.get('/oauth/status/:flowId', async (req, res) => {
|
||||
router.get('/oauth/status/:flowId', requireJwtAuth, async (req, res) => {
|
||||
try {
|
||||
const { flowId } = req.params;
|
||||
const user = req.user;
|
||||
|
||||
if (!user?.id) {
|
||||
return res.status(401).json({ error: 'User not authenticated' });
|
||||
}
|
||||
|
||||
if (!flowId.startsWith(`${user.id}:`) && !flowId.startsWith('system:')) {
|
||||
return res.status(403).json({ error: 'Access denied' });
|
||||
}
|
||||
|
||||
const flowsCache = getLogStores(CacheKeys.FLOWS);
|
||||
const flowManager = getFlowStateManager(flowsCache);
|
||||
|
||||
|
|
@ -375,7 +436,7 @@ router.post('/oauth/cancel/:serverName', requireJwtAuth, async (req, res) => {
|
|||
* Reinitialize MCP server
|
||||
* This endpoint allows reinitializing a specific MCP server
|
||||
*/
|
||||
router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
|
||||
router.post('/:serverName/reinitialize', requireJwtAuth, setOAuthSession, async (req, res) => {
|
||||
try {
|
||||
const { serverName } = req.params;
|
||||
const user = createSafeUser(req.user);
|
||||
|
|
@ -421,6 +482,11 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
|
|||
|
||||
const { success, message, oauthRequired, oauthUrl } = result;
|
||||
|
||||
if (oauthRequired) {
|
||||
const flowId = MCPOAuthHandler.generateFlowId(user.id, serverName);
|
||||
setOAuthCsrfCookie(res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
}
|
||||
|
||||
res.json({
|
||||
success,
|
||||
message,
|
||||
|
|
|
|||
|
|
@ -4,10 +4,9 @@ const passport = require('passport');
|
|||
const { randomState } = require('openid-client');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { ErrorTypes } = require('librechat-data-provider');
|
||||
const { isEnabled, createSetBalanceConfig } = require('@librechat/api');
|
||||
const { checkDomainAllowed, loginLimiter, logHeaders, checkBan } = require('~/server/middleware');
|
||||
const { syncUserEntraGroupMemberships } = require('~/server/services/PermissionService');
|
||||
const { setAuthTokens, setOpenIDAuthTokens } = require('~/server/services/AuthService');
|
||||
const { createSetBalanceConfig } = require('@librechat/api');
|
||||
const { checkDomainAllowed, loginLimiter, logHeaders } = require('~/server/middleware');
|
||||
const { createOAuthHandler } = require('~/server/controllers/auth/oauth');
|
||||
const { getAppConfig } = require('~/server/services/Config');
|
||||
const { Balance } = require('~/db/models');
|
||||
|
||||
|
|
@ -26,36 +25,11 @@ const domains = {
|
|||
router.use(logHeaders);
|
||||
router.use(loginLimiter);
|
||||
|
||||
const oauthHandler = async (req, res, next) => {
|
||||
try {
|
||||
if (res.headersSent) {
|
||||
return;
|
||||
}
|
||||
|
||||
await checkBan(req, res);
|
||||
if (req.banned) {
|
||||
return;
|
||||
}
|
||||
if (
|
||||
req.user &&
|
||||
req.user.provider == 'openid' &&
|
||||
isEnabled(process.env.OPENID_REUSE_TOKENS) === true
|
||||
) {
|
||||
await syncUserEntraGroupMemberships(req.user, req.user.tokenset.access_token);
|
||||
setOpenIDAuthTokens(req.user.tokenset, req, res, req.user._id.toString());
|
||||
} else {
|
||||
await setAuthTokens(req.user._id, res);
|
||||
}
|
||||
res.redirect(domains.client);
|
||||
} catch (err) {
|
||||
logger.error('Error in setting authentication tokens:', err);
|
||||
next(err);
|
||||
}
|
||||
};
|
||||
const oauthHandler = createOAuthHandler();
|
||||
|
||||
router.get('/error', (req, res) => {
|
||||
/** A single error message is pushed by passport when authentication fails. */
|
||||
const errorMessage = req.session?.messages?.pop() || 'Unknown error';
|
||||
const errorMessage = req.session?.messages?.pop() || 'Unknown OAuth error';
|
||||
logger.error('Error in OAuth authentication:', {
|
||||
message: errorMessage,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -6,9 +6,10 @@ const {
|
|||
agentPermissionsSchema,
|
||||
promptPermissionsSchema,
|
||||
memoryPermissionsSchema,
|
||||
mcpServersPermissionsSchema,
|
||||
marketplacePermissionsSchema,
|
||||
peoplePickerPermissionsSchema,
|
||||
mcpServersPermissionsSchema,
|
||||
remoteAgentsPermissionsSchema,
|
||||
} = require('librechat-data-provider');
|
||||
const { checkAdmin, requireJwtAuth } = require('~/server/middleware');
|
||||
const { updateRoleByName, getRoleByName } = require('~/models/Role');
|
||||
|
|
@ -51,6 +52,11 @@ const permissionConfigs = {
|
|||
permissionType: PermissionTypes.MARKETPLACE,
|
||||
errorMessage: 'Invalid marketplace permissions.',
|
||||
},
|
||||
'remote-agents': {
|
||||
schema: remoteAgentsPermissionsSchema,
|
||||
permissionType: PermissionTypes.REMOTE_AGENTS,
|
||||
errorMessage: 'Invalid remote agents permissions.',
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -160,4 +166,10 @@ router.put('/:roleName/mcp-servers', checkAdmin, createPermissionUpdateHandler('
|
|||
*/
|
||||
router.put('/:roleName/marketplace', checkAdmin, createPermissionUpdateHandler('marketplace'));
|
||||
|
||||
/**
|
||||
* PUT /api/roles/:roleName/remote-agents
|
||||
* Update remote agents (API) permissions for a specific role
|
||||
*/
|
||||
router.put('/:roleName/remote-agents', checkAdmin, createPermissionUpdateHandler('remote-agents'));
|
||||
|
||||
module.exports = router;
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ const {
|
|||
logAxiosError,
|
||||
refreshAccessToken,
|
||||
GenerationJobManager,
|
||||
createSSRFSafeAgents,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
Time,
|
||||
|
|
@ -133,6 +134,7 @@ async function loadActionSets(searchParams) {
|
|||
* @param {import('zod').ZodTypeAny | undefined} [params.zodSchema] - The Zod schema for tool input validation/definition
|
||||
* @param {{ oauth_client_id?: string; oauth_client_secret?: string; }} params.encrypted - The encrypted values for the action.
|
||||
* @param {string | null} [params.streamId] - The stream ID for resumable streams.
|
||||
* @param {boolean} [params.useSSRFProtection] - When true, uses SSRF-safe HTTP agents that validate resolved IPs at connect time.
|
||||
* @returns { Promise<typeof tool | { _call: (toolInput: Object | string) => unknown}> } An object with `_call` method to execute the tool input.
|
||||
*/
|
||||
async function createActionTool({
|
||||
|
|
@ -145,7 +147,9 @@ async function createActionTool({
|
|||
description,
|
||||
encrypted,
|
||||
streamId = null,
|
||||
useSSRFProtection = false,
|
||||
}) {
|
||||
const ssrfAgents = useSSRFProtection ? createSSRFSafeAgents() : undefined;
|
||||
/** @type {(toolInput: Object | string, config: GraphRunnableConfig) => Promise<unknown>} */
|
||||
const _call = async (toolInput, config) => {
|
||||
try {
|
||||
|
|
@ -201,7 +205,7 @@ async function createActionTool({
|
|||
async () => {
|
||||
const eventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data };
|
||||
if (streamId) {
|
||||
GenerationJobManager.emitChunk(streamId, eventData);
|
||||
await GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
sendEvent(res, eventData);
|
||||
}
|
||||
|
|
@ -231,7 +235,7 @@ async function createActionTool({
|
|||
data.delta.expires_at = undefined;
|
||||
const successEventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data };
|
||||
if (streamId) {
|
||||
GenerationJobManager.emitChunk(streamId, successEventData);
|
||||
await GenerationJobManager.emitChunk(streamId, successEventData);
|
||||
} else {
|
||||
sendEvent(res, successEventData);
|
||||
}
|
||||
|
|
@ -324,7 +328,7 @@ async function createActionTool({
|
|||
}
|
||||
}
|
||||
|
||||
const response = await preparedExecutor.execute();
|
||||
const response = await preparedExecutor.execute(ssrfAgents);
|
||||
|
||||
if (typeof response.data === 'object') {
|
||||
return JSON.stringify(response.data);
|
||||
|
|
|
|||
|
|
@ -73,15 +73,25 @@ const replaceArtifactContent = (originalText, artifact, original, updated) => {
|
|||
return null;
|
||||
}
|
||||
|
||||
// Check if there are code blocks
|
||||
const codeBlockStart = artifactContent.indexOf('```\n', contentStart);
|
||||
// Check if there are code blocks - handle both ```\n and ```lang\n formats
|
||||
let codeBlockStart = artifactContent.indexOf('```', contentStart);
|
||||
const codeBlockEnd = artifactContent.lastIndexOf('\n```', contentEnd);
|
||||
|
||||
// If we found opening backticks, find the actual newline (skipping any language identifier)
|
||||
if (codeBlockStart !== -1) {
|
||||
const newlineAfterBackticks = artifactContent.indexOf('\n', codeBlockStart);
|
||||
if (newlineAfterBackticks !== -1 && newlineAfterBackticks < contentEnd) {
|
||||
codeBlockStart = newlineAfterBackticks;
|
||||
} else {
|
||||
codeBlockStart = -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Determine where to look for the original content
|
||||
let searchStart, searchEnd;
|
||||
if (codeBlockStart !== -1) {
|
||||
// Code block starts
|
||||
searchStart = codeBlockStart + 4; // after ```\n
|
||||
// Code block starts - searchStart is right after the newline following ```[lang]
|
||||
searchStart = codeBlockStart + 1; // after the newline
|
||||
|
||||
if (codeBlockEnd !== -1 && codeBlockEnd > codeBlockStart) {
|
||||
// Code block has proper ending
|
||||
|
|
|
|||
|
|
@ -494,5 +494,268 @@ ${original}`;
|
|||
/```\n {2}function test\(\) \{\n {4}return \{\n {6}value: 100\n {4}\};\n {2}\}\n```/,
|
||||
);
|
||||
});
|
||||
|
||||
test('should handle code blocks with language identifiers (```svg, ```html, etc.)', () => {
|
||||
const svgContent = `<svg viewBox="0 0 200 200" xmlns="http://www.w3.org/2000/svg">
|
||||
<rect width="200" height="200" fill="#4A90A4"/>
|
||||
<rect x="50" y="50" width="100" height="100" fill="#FFFFFF"/>
|
||||
</svg>`;
|
||||
|
||||
/** Artifact with language identifier in code block */
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test-svg" type="image/svg+xml" title="Test SVG"}
|
||||
\`\`\`svg
|
||||
${svgContent}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
expect(artifacts).toHaveLength(1);
|
||||
|
||||
const updatedSvg = svgContent.replace('#FFFFFF', '#131313');
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], svgContent, updatedSvg);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('#131313');
|
||||
expect(result).not.toContain('#FFFFFF');
|
||||
expect(result).toMatch(/```svg\n/);
|
||||
});
|
||||
|
||||
test('should handle code blocks with complex language identifiers', () => {
|
||||
const htmlContent = `<!DOCTYPE html>
|
||||
<html>
|
||||
<head><title>Test</title></head>
|
||||
<body>Hello</body>
|
||||
</html>`;
|
||||
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test-html" type="text/html" title="Test HTML"}
|
||||
\`\`\`html
|
||||
${htmlContent}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const updatedHtml = htmlContent.replace('Hello', 'Updated');
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], htmlContent, updatedHtml);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('Updated');
|
||||
expect(result).toMatch(/```html\n/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('code block edge cases', () => {
|
||||
test('should handle code block without language identifier (```\\n)', () => {
|
||||
const content = 'const x = 1;\nconst y = 2;';
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test" type="text/plain" title="Test"}
|
||||
\`\`\`
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, 'updated');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated');
|
||||
expect(result).toMatch(/```\nupdated\n```/);
|
||||
});
|
||||
|
||||
test('should handle various language identifiers', () => {
|
||||
const languages = [
|
||||
'javascript',
|
||||
'typescript',
|
||||
'python',
|
||||
'jsx',
|
||||
'tsx',
|
||||
'css',
|
||||
'json',
|
||||
'xml',
|
||||
'markdown',
|
||||
'md',
|
||||
];
|
||||
|
||||
for (const lang of languages) {
|
||||
const content = `test content for ${lang}`;
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test-${lang}" type="text/plain" title="Test"}
|
||||
\`\`\`${lang}
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
expect(artifacts).toHaveLength(1);
|
||||
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, 'updated');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated');
|
||||
expect(result).toMatch(new RegExp(`\`\`\`${lang}\\n`));
|
||||
}
|
||||
});
|
||||
|
||||
test('should handle single character language identifier', () => {
|
||||
const content = 'single char lang';
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test" type="text/plain" title="Test"}
|
||||
\`\`\`r
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, 'updated');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated');
|
||||
expect(result).toMatch(/```r\n/);
|
||||
});
|
||||
|
||||
test('should handle code block with content that looks like code fence', () => {
|
||||
const content = 'Line 1\nSome text with ``` backticks in middle\nLine 3';
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test" type="text/plain" title="Test"}
|
||||
\`\`\`text
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, 'updated');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated');
|
||||
});
|
||||
|
||||
test('should handle code block with trailing whitespace in language line', () => {
|
||||
const content = 'whitespace test';
|
||||
/** Note: trailing spaces after 'python' */
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test" type="text/plain" title="Test"}
|
||||
\`\`\`python
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, 'updated');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated');
|
||||
});
|
||||
|
||||
test('should handle react/jsx content with complex syntax', () => {
|
||||
const jsxContent = `function App() {
|
||||
const [count, setCount] = useState(0);
|
||||
return (
|
||||
<div className="app">
|
||||
<h1>Count: {count}</h1>
|
||||
<button onClick={() => setCount(c => c + 1)}>
|
||||
Increment
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}`;
|
||||
|
||||
const artifactText = `${ARTIFACT_START}{identifier="react-app" type="application/vnd.react" title="React App"}
|
||||
\`\`\`jsx
|
||||
${jsxContent}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const updatedJsx = jsxContent.replace('Increment', 'Click me');
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], jsxContent, updatedJsx);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('Click me');
|
||||
expect(result).not.toContain('Increment');
|
||||
expect(result).toMatch(/```jsx\n/);
|
||||
});
|
||||
|
||||
test('should handle mermaid diagram content', () => {
|
||||
const mermaidContent = `graph TD
|
||||
A[Start] --> B{Is it?}
|
||||
B -->|Yes| C[OK]
|
||||
B -->|No| D[End]`;
|
||||
|
||||
const artifactText = `${ARTIFACT_START}{identifier="diagram" type="application/vnd.mermaid" title="Flow"}
|
||||
\`\`\`mermaid
|
||||
${mermaidContent}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const updatedMermaid = mermaidContent.replace('Start', 'Begin');
|
||||
const result = replaceArtifactContent(
|
||||
artifactText,
|
||||
artifacts[0],
|
||||
mermaidContent,
|
||||
updatedMermaid,
|
||||
);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('Begin');
|
||||
expect(result).toMatch(/```mermaid\n/);
|
||||
});
|
||||
|
||||
test('should handle artifact without code block (plain text)', () => {
|
||||
const content = 'Just plain text without code fences';
|
||||
const artifactText = `${ARTIFACT_START}{identifier="plain" type="text/plain" title="Plain"}
|
||||
${content}
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const result = replaceArtifactContent(
|
||||
artifactText,
|
||||
artifacts[0],
|
||||
content,
|
||||
'updated plain text',
|
||||
);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('updated plain text');
|
||||
expect(result).not.toContain('```');
|
||||
});
|
||||
|
||||
test('should handle multiline content with various newline patterns', () => {
|
||||
const content = `Line 1
|
||||
Line 2
|
||||
|
||||
Line 4 after empty line
|
||||
Indented line
|
||||
Double indented`;
|
||||
|
||||
const artifactText = `${ARTIFACT_START}{identifier="test" type="text/plain" title="Test"}
|
||||
\`\`\`
|
||||
${content}
|
||||
\`\`\`
|
||||
${ARTIFACT_END}`;
|
||||
|
||||
const message = { text: artifactText };
|
||||
const artifacts = findAllArtifacts(message);
|
||||
|
||||
const updated = content.replace('Line 1', 'First Line');
|
||||
const result = replaceArtifactContent(artifactText, artifacts[0], content, updated);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result).toContain('First Line');
|
||||
expect(result).toContain(' Indented line');
|
||||
expect(result).toContain(' Double indented');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -7,7 +7,13 @@ const {
|
|||
DEFAULT_REFRESH_TOKEN_EXPIRY,
|
||||
} = require('@librechat/data-schemas');
|
||||
const { ErrorTypes, SystemRoles, errorsToString } = require('librechat-data-provider');
|
||||
const { isEnabled, checkEmailConfig, isEmailDomainAllowed, math } = require('@librechat/api');
|
||||
const {
|
||||
math,
|
||||
isEnabled,
|
||||
checkEmailConfig,
|
||||
isEmailDomainAllowed,
|
||||
shouldUseSecureCookie,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
findUser,
|
||||
findToken,
|
||||
|
|
@ -33,7 +39,6 @@ const domains = {
|
|||
server: process.env.DOMAIN_SERVER,
|
||||
};
|
||||
|
||||
const isProduction = process.env.NODE_ENV === 'production';
|
||||
const genericVerificationMessage = 'Please check your email to verify your email address.';
|
||||
|
||||
/**
|
||||
|
|
@ -392,13 +397,13 @@ const setAuthTokens = async (userId, res, _session = null) => {
|
|||
res.cookie('refreshToken', refreshToken, {
|
||||
expires: new Date(refreshTokenExpires),
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
res.cookie('token_provider', 'librechat', {
|
||||
expires: new Date(refreshTokenExpires),
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
return token;
|
||||
|
|
@ -419,7 +424,7 @@ const setAuthTokens = async (userId, res, _session = null) => {
|
|||
* @param {Object} req - request object (for session access)
|
||||
* @param {Object} res - response object
|
||||
* @param {string} [userId] - Optional MongoDB user ID for image path validation
|
||||
* @returns {String} - access token
|
||||
* @returns {String} - id_token (preferred) or access_token as the app auth token
|
||||
*/
|
||||
const setOpenIDAuthTokens = (tokenset, req, res, userId, existingRefreshToken) => {
|
||||
try {
|
||||
|
|
@ -448,34 +453,62 @@ const setOpenIDAuthTokens = (tokenset, req, res, userId, existingRefreshToken) =
|
|||
return;
|
||||
}
|
||||
|
||||
/**
|
||||
* Use id_token as the app authentication token (Bearer token for JWKS validation).
|
||||
* The id_token is always a standard JWT signed by the IdP's JWKS keys with the app's
|
||||
* client_id as audience. The access_token may be opaque or intended for a different
|
||||
* audience (e.g., Microsoft Graph API), which fails JWKS validation.
|
||||
* Falls back to access_token for providers where id_token is not available.
|
||||
*/
|
||||
const appAuthToken = tokenset.id_token || tokenset.access_token;
|
||||
|
||||
/**
|
||||
* Always set refresh token cookie so it survives express session expiry.
|
||||
* The session cookie maxAge (SESSION_EXPIRY, default 15 min) is typically shorter
|
||||
* than the OIDC token lifetime (~1 hour). Without this cookie fallback, the refresh
|
||||
* token stored only in the session is lost when the session expires, causing the user
|
||||
* to be signed out on the next token refresh attempt.
|
||||
* The refresh token is small (opaque string) so it doesn't hit the HTTP/2 header
|
||||
* size limits that motivated session storage for the larger access_token/id_token.
|
||||
*/
|
||||
res.cookie('refreshToken', refreshToken, {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
|
||||
/** Store tokens server-side in session to avoid large cookies */
|
||||
if (req.session) {
|
||||
req.session.openidTokens = {
|
||||
accessToken: tokenset.access_token,
|
||||
idToken: tokenset.id_token,
|
||||
refreshToken: refreshToken,
|
||||
expiresAt: expirationDate.getTime(),
|
||||
};
|
||||
} else {
|
||||
logger.warn('[setOpenIDAuthTokens] No session available, falling back to cookies');
|
||||
res.cookie('refreshToken', refreshToken, {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
sameSite: 'strict',
|
||||
});
|
||||
res.cookie('openid_access_token', tokenset.access_token, {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
if (tokenset.id_token) {
|
||||
res.cookie('openid_id_token', tokenset.id_token, {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/** Small cookie to indicate token provider (required for auth middleware) */
|
||||
res.cookie('token_provider', 'openid', {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
if (userId && isEnabled(process.env.OPENID_REUSE_TOKENS)) {
|
||||
|
|
@ -486,11 +519,11 @@ const setOpenIDAuthTokens = (tokenset, req, res, userId, existingRefreshToken) =
|
|||
res.cookie('openid_user_id', signedUserId, {
|
||||
expires: expirationDate,
|
||||
httpOnly: true,
|
||||
secure: isProduction,
|
||||
secure: shouldUseSecureCookie(),
|
||||
sameSite: 'strict',
|
||||
});
|
||||
}
|
||||
return tokenset.access_token;
|
||||
return appAuthToken;
|
||||
} catch (error) {
|
||||
logger.error('[setOpenIDAuthTokens] Error in setting authentication tokens:', error);
|
||||
throw error;
|
||||
|
|
|
|||
269
api/server/services/AuthService.spec.js
Normal file
269
api/server/services/AuthService.spec.js
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: { info: jest.fn(), warn: jest.fn(), debug: jest.fn(), error: jest.fn() },
|
||||
DEFAULT_SESSION_EXPIRY: 900000,
|
||||
DEFAULT_REFRESH_TOKEN_EXPIRY: 604800000,
|
||||
}));
|
||||
jest.mock('librechat-data-provider', () => ({
|
||||
ErrorTypes: {},
|
||||
SystemRoles: { USER: 'USER', ADMIN: 'ADMIN' },
|
||||
errorsToString: jest.fn(),
|
||||
}));
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: jest.fn((val) => val === 'true' || val === true),
|
||||
checkEmailConfig: jest.fn(),
|
||||
isEmailDomainAllowed: jest.fn(),
|
||||
math: jest.fn((val, fallback) => (val ? Number(val) : fallback)),
|
||||
shouldUseSecureCookie: jest.fn(() => false),
|
||||
}));
|
||||
jest.mock('~/models', () => ({
|
||||
findUser: jest.fn(),
|
||||
findToken: jest.fn(),
|
||||
createUser: jest.fn(),
|
||||
updateUser: jest.fn(),
|
||||
countUsers: jest.fn(),
|
||||
getUserById: jest.fn(),
|
||||
findSession: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
deleteSession: jest.fn(),
|
||||
createSession: jest.fn(),
|
||||
generateToken: jest.fn(),
|
||||
deleteUserById: jest.fn(),
|
||||
generateRefreshToken: jest.fn(),
|
||||
}));
|
||||
jest.mock('~/strategies/validators', () => ({ registerSchema: { parse: jest.fn() } }));
|
||||
jest.mock('~/server/services/Config', () => ({ getAppConfig: jest.fn() }));
|
||||
jest.mock('~/server/utils', () => ({ sendEmail: jest.fn() }));
|
||||
|
||||
const { shouldUseSecureCookie } = require('@librechat/api');
|
||||
const { setOpenIDAuthTokens } = require('./AuthService');
|
||||
|
||||
/** Helper to build a mock Express response */
|
||||
function mockResponse() {
|
||||
const cookies = {};
|
||||
const res = {
|
||||
cookie: jest.fn((name, value, options) => {
|
||||
cookies[name] = { value, options };
|
||||
}),
|
||||
_cookies: cookies,
|
||||
};
|
||||
return res;
|
||||
}
|
||||
|
||||
/** Helper to build a mock Express request with session */
|
||||
function mockRequest(sessionData = {}) {
|
||||
return {
|
||||
session: { openidTokens: null, ...sessionData },
|
||||
};
|
||||
}
|
||||
|
||||
describe('setOpenIDAuthTokens', () => {
|
||||
const env = process.env;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
process.env = {
|
||||
...env,
|
||||
JWT_REFRESH_SECRET: 'test-refresh-secret',
|
||||
OPENID_REUSE_TOKENS: 'true',
|
||||
};
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
process.env = env;
|
||||
});
|
||||
|
||||
describe('token selection (id_token vs access_token)', () => {
|
||||
it('should return id_token when both id_token and access_token are present', () => {
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBe('the-id-token');
|
||||
});
|
||||
|
||||
it('should return access_token when id_token is not available', () => {
|
||||
const tokenset = {
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBe('the-access-token');
|
||||
});
|
||||
|
||||
it('should return access_token when id_token is undefined', () => {
|
||||
const tokenset = {
|
||||
id_token: undefined,
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBe('the-access-token');
|
||||
});
|
||||
|
||||
it('should return access_token when id_token is null', () => {
|
||||
const tokenset = {
|
||||
id_token: null,
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBe('the-access-token');
|
||||
});
|
||||
|
||||
it('should return id_token even when id_token and access_token differ', () => {
|
||||
const tokenset = {
|
||||
id_token: 'id-token-jwt-signed-by-idp',
|
||||
access_token: 'opaque-graph-api-token',
|
||||
refresh_token: 'refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBe('id-token-jwt-signed-by-idp');
|
||||
expect(result).not.toBe('opaque-graph-api-token');
|
||||
});
|
||||
});
|
||||
|
||||
describe('session token storage', () => {
|
||||
it('should store the original access_token in session (not id_token)', () => {
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
|
||||
expect(req.session.openidTokens.accessToken).toBe('the-access-token');
|
||||
expect(req.session.openidTokens.refreshToken).toBe('the-refresh-token');
|
||||
});
|
||||
});
|
||||
|
||||
describe('cookie secure flag', () => {
|
||||
it('should call shouldUseSecureCookie for every cookie set', () => {
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
|
||||
// token_provider + openid_user_id (session path, so no refreshToken/openid_access_token cookies)
|
||||
const secureCalls = shouldUseSecureCookie.mock.calls.length;
|
||||
expect(secureCalls).toBeGreaterThanOrEqual(2);
|
||||
|
||||
// Verify all cookies use the result of shouldUseSecureCookie
|
||||
for (const [, cookie] of Object.entries(res._cookies)) {
|
||||
expect(cookie.options.secure).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it('should set secure: true when shouldUseSecureCookie returns true', () => {
|
||||
shouldUseSecureCookie.mockReturnValue(true);
|
||||
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
|
||||
for (const [, cookie] of Object.entries(res._cookies)) {
|
||||
expect(cookie.options.secure).toBe(true);
|
||||
}
|
||||
});
|
||||
|
||||
it('should use shouldUseSecureCookie for cookie fallback path (no session)', () => {
|
||||
shouldUseSecureCookie.mockReturnValue(false);
|
||||
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
refresh_token: 'the-refresh-token',
|
||||
};
|
||||
const req = { session: null };
|
||||
const res = mockResponse();
|
||||
|
||||
setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
|
||||
// In the cookie fallback path, we get: refreshToken, openid_access_token, token_provider, openid_user_id
|
||||
expect(res.cookie).toHaveBeenCalledWith(
|
||||
'refreshToken',
|
||||
expect.any(String),
|
||||
expect.objectContaining({ secure: false }),
|
||||
);
|
||||
expect(res.cookie).toHaveBeenCalledWith(
|
||||
'openid_access_token',
|
||||
expect.any(String),
|
||||
expect.objectContaining({ secure: false }),
|
||||
);
|
||||
expect(res.cookie).toHaveBeenCalledWith(
|
||||
'token_provider',
|
||||
'openid',
|
||||
expect.objectContaining({ secure: false }),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('should return undefined when tokenset is null', () => {
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
const result = setOpenIDAuthTokens(null, req, res, 'user-123');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return undefined when access_token is missing', () => {
|
||||
const tokenset = { refresh_token: 'refresh' };
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return undefined when no refresh token is available', () => {
|
||||
const tokenset = { access_token: 'access', id_token: 'id' };
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should use existingRefreshToken when tokenset has no refresh_token', () => {
|
||||
const tokenset = {
|
||||
id_token: 'the-id-token',
|
||||
access_token: 'the-access-token',
|
||||
};
|
||||
const req = mockRequest();
|
||||
const res = mockResponse();
|
||||
|
||||
const result = setOpenIDAuthTokens(tokenset, req, res, 'user-123', 'existing-refresh');
|
||||
expect(result).toBe('the-id-token');
|
||||
expect(req.session.openidTokens.refreshToken).toBe('existing-refresh');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,10 +1,92 @@
|
|||
const { ToolCacheKeys } = require('../getCachedTools');
|
||||
const { CacheKeys } = require('librechat-data-provider');
|
||||
|
||||
jest.mock('~/cache/getLogStores');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
|
||||
const mockCache = { get: jest.fn(), set: jest.fn(), delete: jest.fn() };
|
||||
getLogStores.mockReturnValue(mockCache);
|
||||
|
||||
const {
|
||||
ToolCacheKeys,
|
||||
getCachedTools,
|
||||
setCachedTools,
|
||||
getMCPServerTools,
|
||||
invalidateCachedTools,
|
||||
} = require('../getCachedTools');
|
||||
|
||||
describe('getCachedTools', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
getLogStores.mockReturnValue(mockCache);
|
||||
});
|
||||
|
||||
describe('getCachedTools - Cache Isolation Security', () => {
|
||||
describe('ToolCacheKeys.MCP_SERVER', () => {
|
||||
it('should generate cache keys that include userId', () => {
|
||||
const key = ToolCacheKeys.MCP_SERVER('user123', 'github');
|
||||
expect(key).toBe('tools:mcp:user123:github');
|
||||
});
|
||||
});
|
||||
|
||||
describe('TOOL_CACHE namespace usage', () => {
|
||||
it('getCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
await getCachedTools();
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
});
|
||||
|
||||
it('getCachedTools with MCP server options should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue({ tool1: {} });
|
||||
await getCachedTools({ userId: 'user1', serverName: 'github' });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
});
|
||||
|
||||
it('setCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
const tools = { tool1: { type: 'function' } };
|
||||
await setCachedTools(tools);
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL, tools, expect.any(Number));
|
||||
});
|
||||
|
||||
it('setCachedTools with MCP server options should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
const tools = { tool1: { type: 'function' } };
|
||||
await setCachedTools(tools, { userId: 'user1', serverName: 'github' });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github'),
|
||||
tools,
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('invalidateCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL);
|
||||
});
|
||||
|
||||
it('getMCPServerTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
await getMCPServerTools('user1', 'github');
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
});
|
||||
|
||||
it('should NOT use CONFIG_STORE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
await getCachedTools();
|
||||
await getMCPServerTools('user1', 'github');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
await setCachedTools({ tool1: {} });
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
|
||||
const allCalls = getLogStores.mock.calls.flat();
|
||||
expect(allCalls).not.toContain(CacheKeys.CONFIG_STORE);
|
||||
expect(allCalls.every((key) => key === CacheKeys.TOOL_CACHE)).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ const ToolCacheKeys = {
|
|||
* @returns {Promise<LCAvailableTools|null>} The available tools object or null if not cached
|
||||
*/
|
||||
async function getCachedTools(options = {}) {
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName } = options;
|
||||
|
||||
// Return MCP server-specific tools if requested
|
||||
|
|
@ -43,7 +43,7 @@ async function getCachedTools(options = {}) {
|
|||
* @returns {Promise<boolean>} Whether the operation was successful
|
||||
*/
|
||||
async function setCachedTools(tools, options = {}) {
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName, ttl = Time.TWELVE_HOURS } = options;
|
||||
|
||||
// Cache by MCP server if specified (requires userId)
|
||||
|
|
@ -65,7 +65,7 @@ async function setCachedTools(tools, options = {}) {
|
|||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function invalidateCachedTools(options = {}) {
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName, invalidateGlobal = false } = options;
|
||||
|
||||
const keysToDelete = [];
|
||||
|
|
@ -89,7 +89,7 @@ async function invalidateCachedTools(options = {}) {
|
|||
* @returns {Promise<LCAvailableTools|null>} The available tools for the server
|
||||
*/
|
||||
async function getMCPServerTools(userId, serverName) {
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const serverTools = await cache.get(ToolCacheKeys.MCP_SERVER(userId, serverName));
|
||||
|
||||
if (serverTools) {
|
||||
|
|
|
|||
|
|
@ -28,6 +28,11 @@ async function loadConfigModels(req) {
|
|||
modelsConfig[EModelEndpoint.azureAssistants] = azureConfig.assistantModels;
|
||||
}
|
||||
|
||||
const bedrockConfig = appConfig.endpoints?.[EModelEndpoint.bedrock];
|
||||
if (bedrockConfig?.models && Array.isArray(bedrockConfig.models)) {
|
||||
modelsConfig[EModelEndpoint.bedrock] = bedrockConfig.models;
|
||||
}
|
||||
|
||||
if (!Array.isArray(appConfig.endpoints?.[EModelEndpoint.custom])) {
|
||||
return modelsConfig;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ async function updateMCPServerTools({ userId, serverName, tools }) {
|
|||
|
||||
await setCachedTools(serverTools, { userId, serverName });
|
||||
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
await cache.delete(CacheKeys.TOOLS);
|
||||
logger.debug(
|
||||
`[MCP Cache] Updated ${tools.length} tools for server ${serverName} (user: ${userId})`,
|
||||
|
|
@ -61,7 +61,7 @@ async function mergeAppTools(appTools) {
|
|||
const cachedTools = await getCachedTools();
|
||||
const mergedTools = { ...cachedTools, ...appTools };
|
||||
await setCachedTools(mergedTools);
|
||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
await cache.delete(CacheKeys.TOOLS);
|
||||
logger.debug(`Merged ${count} app-level tools`);
|
||||
} catch (error) {
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ setGetAgent(getAgent);
|
|||
* @param {Function} params.loadTools - Function to load agent tools
|
||||
* @param {Array} params.requestFiles - Request files
|
||||
* @param {string} params.conversationId - The conversation ID
|
||||
* @param {string} [params.parentMessageId] - The parent message ID for thread filtering
|
||||
* @param {Set} params.allowedProviders - Set of allowed providers
|
||||
* @param {Map} params.agentConfigs - Map of agent configs to add to
|
||||
* @param {string} params.primaryAgentId - The primary agent ID
|
||||
|
|
@ -46,6 +47,7 @@ const processAddedConvo = async ({
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
allowedProviders,
|
||||
agentConfigs,
|
||||
primaryAgentId,
|
||||
|
|
@ -91,6 +93,7 @@ const processAddedConvo = async ({
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent: addedAgent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
|
|
@ -99,9 +102,12 @@ const processAddedConvo = async ({
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ const { createContentAggregator } = require('@librechat/agents');
|
|||
const {
|
||||
initializeAgent,
|
||||
validateAgentModel,
|
||||
getCustomEndpointConfig,
|
||||
createSequentialChainEdges,
|
||||
createEdgeCollector,
|
||||
filterOrphanedEdges,
|
||||
GenerationJobManager,
|
||||
getCustomEndpointConfig,
|
||||
createSequentialChainEdges,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
EModelEndpoint,
|
||||
|
|
@ -18,8 +19,8 @@ const {
|
|||
createToolEndCallback,
|
||||
getDefaultHandlers,
|
||||
} = require('~/server/controllers/agents/callbacks');
|
||||
const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
|
||||
const { getModelsConfig } = require('~/server/controllers/ModelController');
|
||||
const { loadAgentTools } = require('~/server/services/ToolService');
|
||||
const AgentClient = require('~/server/controllers/agents/client');
|
||||
const { getConvoFiles } = require('~/models/Conversation');
|
||||
const { processAddedConvo } = require('./addedConvo');
|
||||
|
|
@ -31,8 +32,10 @@ const db = require('~/models');
|
|||
* Creates a tool loader function for the agent.
|
||||
* @param {AbortSignal} signal - The abort signal
|
||||
* @param {string | null} [streamId] - The stream ID for resumable mode
|
||||
* @param {boolean} [definitionsOnly=false] - When true, returns only serializable
|
||||
* tool definitions without creating full tool instances (for event-driven mode)
|
||||
*/
|
||||
function createToolLoader(signal, streamId = null) {
|
||||
function createToolLoader(signal, streamId = null, definitionsOnly = false) {
|
||||
/**
|
||||
* @param {object} params
|
||||
* @param {ServerRequest} params.req
|
||||
|
|
@ -43,21 +46,33 @@ function createToolLoader(signal, streamId = null) {
|
|||
* @param {string} params.model
|
||||
* @param {AgentToolResources} params.tool_resources
|
||||
* @returns {Promise<{
|
||||
* tools: StructuredTool[],
|
||||
* toolContextMap: Record<string, unknown>,
|
||||
* userMCPAuthMap?: Record<string, Record<string, string>>
|
||||
* tools?: StructuredTool[],
|
||||
* toolContextMap: Record<string, unknown>,
|
||||
* toolDefinitions?: import('@librechat/agents').LCTool[],
|
||||
* userMCPAuthMap?: Record<string, Record<string, string>>,
|
||||
* toolRegistry?: import('@librechat/agents').LCToolRegistry
|
||||
* } | undefined>}
|
||||
*/
|
||||
return async function loadTools({ req, res, agentId, tools, provider, model, tool_resources }) {
|
||||
const agent = { id: agentId, tools, provider, model };
|
||||
return async function loadTools({
|
||||
req,
|
||||
res,
|
||||
tools,
|
||||
model,
|
||||
agentId,
|
||||
provider,
|
||||
tool_options,
|
||||
tool_resources,
|
||||
}) {
|
||||
const agent = { id: agentId, tools, provider, model, tool_options };
|
||||
try {
|
||||
return await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent,
|
||||
signal,
|
||||
tool_resources,
|
||||
streamId,
|
||||
tool_resources,
|
||||
definitionsOnly,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('Error loading tools for agent ' + agentId, error);
|
||||
|
|
@ -80,8 +95,47 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
const artifactPromises = [];
|
||||
const { contentParts, aggregateContent } = createContentAggregator();
|
||||
const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId });
|
||||
|
||||
/**
|
||||
* Agent context store - populated after initialization, accessed by callback via closure.
|
||||
* Maps agentId -> { userMCPAuthMap, agent, tool_resources, toolRegistry, openAIApiKey }
|
||||
* @type {Map<string, {
|
||||
* userMCPAuthMap?: Record<string, Record<string, string>>,
|
||||
* agent?: object,
|
||||
* tool_resources?: object,
|
||||
* toolRegistry?: import('@librechat/agents').LCToolRegistry,
|
||||
* openAIApiKey?: string
|
||||
* }>}
|
||||
*/
|
||||
const agentToolContexts = new Map();
|
||||
|
||||
const toolExecuteOptions = {
|
||||
loadTools: async (toolNames, agentId) => {
|
||||
const ctx = agentToolContexts.get(agentId) ?? {};
|
||||
logger.debug(`[ON_TOOL_EXECUTE] ctx found: ${!!ctx.userMCPAuthMap}, agent: ${ctx.agent?.id}`);
|
||||
logger.debug(`[ON_TOOL_EXECUTE] toolRegistry size: ${ctx.toolRegistry?.size ?? 'undefined'}`);
|
||||
|
||||
const result = await loadToolsForExecution({
|
||||
req,
|
||||
res,
|
||||
signal,
|
||||
streamId,
|
||||
toolNames,
|
||||
agent: ctx.agent,
|
||||
toolRegistry: ctx.toolRegistry,
|
||||
userMCPAuthMap: ctx.userMCPAuthMap,
|
||||
tool_resources: ctx.tool_resources,
|
||||
});
|
||||
|
||||
logger.debug(`[ON_TOOL_EXECUTE] loaded ${result.loadedTools?.length ?? 0} tools`);
|
||||
return result;
|
||||
},
|
||||
toolEndCallback,
|
||||
};
|
||||
|
||||
const eventHandlers = getDefaultHandlers({
|
||||
res,
|
||||
toolExecuteOptions,
|
||||
aggregateContent,
|
||||
toolEndCallback,
|
||||
collectedUsage,
|
||||
|
|
@ -114,11 +168,14 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
const agentConfigs = new Map();
|
||||
const allowedProviders = new Set(appConfig?.endpoints?.[EModelEndpoint.agents]?.allowedProviders);
|
||||
|
||||
const loadTools = createToolLoader(signal, streamId);
|
||||
/** Event-driven mode: only load tool definitions, not full instances */
|
||||
const loadTools = createToolLoader(signal, streamId, true);
|
||||
/** @type {Array<MongoFile>} */
|
||||
const requestFiles = req.body.files ?? [];
|
||||
/** @type {string} */
|
||||
const conversationId = req.body.conversationId;
|
||||
/** @type {string | undefined} */
|
||||
const parentMessageId = req.body.parentMessageId;
|
||||
|
||||
const primaryConfig = await initializeAgent(
|
||||
{
|
||||
|
|
@ -127,6 +184,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent: primaryAgent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
|
|
@ -136,12 +194,31 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
logger.debug(
|
||||
`[initializeClient] Tool definitions for primary agent: ${primaryConfig.toolDefinitions?.length ?? 0}`,
|
||||
);
|
||||
|
||||
/** Store primary agent's tool context for ON_TOOL_EXECUTE callback */
|
||||
logger.debug(`[initializeClient] Storing tool context for agentId: ${primaryConfig.id}`);
|
||||
logger.debug(
|
||||
`[initializeClient] toolRegistry size: ${primaryConfig.toolRegistry?.size ?? 'undefined'}`,
|
||||
);
|
||||
agentToolContexts.set(primaryConfig.id, {
|
||||
agent: primaryAgent,
|
||||
toolRegistry: primaryConfig.toolRegistry,
|
||||
userMCPAuthMap: primaryConfig.userMCPAuthMap,
|
||||
tool_resources: primaryConfig.tool_resources,
|
||||
});
|
||||
|
||||
const agent_ids = primaryConfig.agent_ids;
|
||||
let userMCPAuthMap = primaryConfig.userMCPAuthMap;
|
||||
|
||||
|
|
@ -178,6 +255,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
},
|
||||
|
|
@ -185,16 +263,29 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
if (userMCPAuthMap != null) {
|
||||
Object.assign(userMCPAuthMap, config.userMCPAuthMap ?? {});
|
||||
} else {
|
||||
userMCPAuthMap = config.userMCPAuthMap;
|
||||
}
|
||||
|
||||
/** Store handoff agent's tool context for ON_TOOL_EXECUTE callback */
|
||||
agentToolContexts.set(agentId, {
|
||||
agent,
|
||||
toolRegistry: config.toolRegistry,
|
||||
userMCPAuthMap: config.userMCPAuthMap,
|
||||
tool_resources: config.tool_resources,
|
||||
});
|
||||
|
||||
agentConfigs.set(agentId, config);
|
||||
return agent;
|
||||
}
|
||||
|
|
@ -242,17 +333,18 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
const { userMCPAuthMap: updatedMCPAuthMap } = await processAddedConvo({
|
||||
req,
|
||||
res,
|
||||
endpointOption,
|
||||
modelsConfig,
|
||||
logViolation,
|
||||
loadTools,
|
||||
logViolation,
|
||||
modelsConfig,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
allowedProviders,
|
||||
agentConfigs,
|
||||
primaryAgentId: primaryConfig.id,
|
||||
primaryAgent,
|
||||
endpointOption,
|
||||
userMCPAuthMap,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
allowedProviders,
|
||||
primaryAgentId: primaryConfig.id,
|
||||
});
|
||||
|
||||
if (updatedMCPAuthMap) {
|
||||
|
|
@ -314,6 +406,10 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
endpoint: isEphemeralAgentId(primaryConfig.id) ? primaryConfig.endpoint : EModelEndpoint.agents,
|
||||
});
|
||||
|
||||
if (streamId) {
|
||||
GenerationJobManager.setCollectedUsage(streamId, collectedUsage);
|
||||
}
|
||||
|
||||
return { client, userMCPAuthMap };
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ const addTitle = async (req, { text, response, client }) => {
|
|||
conversationId: response.conversationId,
|
||||
title,
|
||||
},
|
||||
{ context: 'api/server/services/Endpoints/agents/title.js' },
|
||||
{ context: 'api/server/services/Endpoints/agents/title.js', noUpsert: true },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('Error generating title:', error);
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ const addTitle = async (req, { text, responseText, conversationId }) => {
|
|||
conversationId,
|
||||
title,
|
||||
},
|
||||
{ context: 'api/server/services/Endpoints/assistants/addTitle.js' },
|
||||
{ context: 'api/server/services/Endpoints/assistants/addTitle.js', noUpsert: true },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[addTitle] Error generating title:', error);
|
||||
|
|
@ -81,7 +81,7 @@ const addTitle = async (req, { text, responseText, conversationId }) => {
|
|||
conversationId,
|
||||
title: fallbackTitle,
|
||||
},
|
||||
{ context: 'api/server/services/Endpoints/assistants/addTitle.js' },
|
||||
{ context: 'api/server/services/Endpoints/assistants/addTitle.js', noUpsert: true },
|
||||
);
|
||||
}
|
||||
};
|
||||
|
|
|
|||
|
|
@ -128,7 +128,6 @@ const initializeClient = async ({ req, res, version, endpointOption, initAppClie
|
|||
const groupName = modelGroupMap[modelName].group;
|
||||
clientOptions.addParams = azureConfig.groupMap[groupName].addParams;
|
||||
clientOptions.dropParams = azureConfig.groupMap[groupName].dropParams;
|
||||
clientOptions.forcePrompt = azureConfig.groupMap[groupName].forcePrompt;
|
||||
|
||||
clientOptions.reverseProxyUrl = baseURL ?? clientOptions.reverseProxyUrl;
|
||||
clientOptions.headers = opts.defaultHeaders;
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ const initGoogle = require('~/server/services/Endpoints/google/initialize');
|
|||
* @returns {boolean} - True if the provider is a known custom provider, false otherwise
|
||||
*/
|
||||
function isKnownCustomProvider(provider) {
|
||||
return [Providers.XAI, Providers.DEEPSEEK, Providers.OPENROUTER].includes(
|
||||
return [Providers.XAI, Providers.DEEPSEEK, Providers.OPENROUTER, Providers.MOONSHOT].includes(
|
||||
provider?.toLowerCase() || '',
|
||||
);
|
||||
}
|
||||
|
|
@ -20,6 +20,7 @@ function isKnownCustomProvider(provider) {
|
|||
const providerConfigMap = {
|
||||
[Providers.XAI]: initCustom,
|
||||
[Providers.DEEPSEEK]: initCustom,
|
||||
[Providers.MOONSHOT]: initCustom,
|
||||
[Providers.OPENROUTER]: initCustom,
|
||||
[EModelEndpoint.openAI]: initOpenAI,
|
||||
[EModelEndpoint.google]: initGoogle,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ const mime = require('mime');
|
|||
const axios = require('axios');
|
||||
const fetch = require('node-fetch');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getAzureContainerClient } = require('@librechat/api');
|
||||
const { getAzureContainerClient, deleteRagFile } = require('@librechat/api');
|
||||
|
||||
const defaultBasePath = 'images';
|
||||
const { AZURE_STORAGE_PUBLIC_ACCESS = 'true', AZURE_CONTAINER_NAME = 'files' } = process.env;
|
||||
|
|
@ -102,6 +102,8 @@ async function getAzureURL({ fileName, basePath = defaultBasePath, userId, conta
|
|||
* @param {MongoFile} params.file - The file object.
|
||||
*/
|
||||
async function deleteFileFromAzure(req, file) {
|
||||
await deleteRagFile({ userId: req.user.id, file });
|
||||
|
||||
try {
|
||||
const containerClient = await getAzureContainerClient(AZURE_CONTAINER_NAME);
|
||||
const blobPath = file.filepath.split(`${AZURE_CONTAINER_NAME}/`)[1];
|
||||
|
|
|
|||
|
|
@ -6,27 +6,68 @@ const { getCodeBaseURL } = require('@librechat/agents');
|
|||
const { logAxiosError, getBasePath } = require('@librechat/api');
|
||||
const {
|
||||
Tools,
|
||||
megabyte,
|
||||
fileConfig,
|
||||
FileContext,
|
||||
FileSources,
|
||||
imageExtRegex,
|
||||
inferMimeType,
|
||||
EToolResources,
|
||||
EModelEndpoint,
|
||||
mergeFileConfig,
|
||||
getEndpointFileConfig,
|
||||
} = require('librechat-data-provider');
|
||||
const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions');
|
||||
const { createFile, getFiles, updateFile, claimCodeFile } = require('~/models');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { convertImage } = require('~/server/services/Files/images/convert');
|
||||
const { createFile, getFiles, updateFile } = require('~/models');
|
||||
const { determineFileType } = require('~/server/utils');
|
||||
|
||||
/**
|
||||
* Process OpenAI image files, convert to target format, save and return file metadata.
|
||||
* Creates a fallback download URL response when file cannot be processed locally.
|
||||
* Used when: file exceeds size limit, storage strategy unavailable, or download error occurs.
|
||||
* @param {Object} params - The parameters.
|
||||
* @param {string} params.name - The filename.
|
||||
* @param {string} params.session_id - The code execution session ID.
|
||||
* @param {string} params.id - The file ID from the code environment.
|
||||
* @param {string} params.conversationId - The current conversation ID.
|
||||
* @param {string} params.toolCallId - The tool call ID that generated the file.
|
||||
* @param {string} params.messageId - The current message ID.
|
||||
* @param {number} params.expiresAt - Expiration timestamp (24 hours from creation).
|
||||
* @returns {Object} Fallback response with download URL.
|
||||
*/
|
||||
const createDownloadFallback = ({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
expiresAt,
|
||||
session_id,
|
||||
toolCallId,
|
||||
conversationId,
|
||||
}) => {
|
||||
const basePath = getBasePath();
|
||||
return {
|
||||
filename: name,
|
||||
filepath: `${basePath}/api/files/code/download/${session_id}/${id}`,
|
||||
expiresAt,
|
||||
conversationId,
|
||||
toolCallId,
|
||||
messageId,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Process code execution output files - downloads and saves both images and non-image files.
|
||||
* All files are saved to local storage with fileIdentifier metadata for code env re-upload.
|
||||
* @param {ServerRequest} params.req - The Express request object.
|
||||
* @param {string} params.id - The file ID.
|
||||
* @param {string} params.id - The file ID from the code environment.
|
||||
* @param {string} params.name - The filename.
|
||||
* @param {string} params.apiKey - The code execution API key.
|
||||
* @param {string} params.toolCallId - The tool call ID that generated the file.
|
||||
* @param {string} params.session_id - The code execution session ID.
|
||||
* @param {string} params.conversationId - The current conversation ID.
|
||||
* @param {string} params.messageId - The current message ID.
|
||||
* @returns {Promise<MongoFile & { messageId: string, toolCallId: string } | { filename: string; filepath: string; expiresAt: number; conversationId: string; toolCallId: string; messageId: string } | undefined>} The file metadata or undefined if an error occurs.
|
||||
* @returns {Promise<MongoFile & { messageId: string, toolCallId: string } | undefined>} The file metadata or undefined if an error occurs.
|
||||
*/
|
||||
const processCodeOutput = async ({
|
||||
req,
|
||||
|
|
@ -41,19 +82,15 @@ const processCodeOutput = async ({
|
|||
const appConfig = req.config;
|
||||
const currentDate = new Date();
|
||||
const baseURL = getCodeBaseURL();
|
||||
const basePath = getBasePath();
|
||||
const fileExt = path.extname(name);
|
||||
if (!fileExt || !imageExtRegex.test(name)) {
|
||||
return {
|
||||
filename: name,
|
||||
filepath: `${basePath}/api/files/code/download/${session_id}/${id}`,
|
||||
/** Note: expires 24 hours after creation */
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
conversationId,
|
||||
toolCallId,
|
||||
messageId,
|
||||
};
|
||||
}
|
||||
const fileExt = path.extname(name).toLowerCase();
|
||||
const isImage = fileExt && imageExtRegex.test(name);
|
||||
|
||||
const mergedFileConfig = mergeFileConfig(appConfig.fileConfig);
|
||||
const endpointFileConfig = getEndpointFileConfig({
|
||||
fileConfig: mergedFileConfig,
|
||||
endpoint: EModelEndpoint.agents,
|
||||
});
|
||||
const fileSizeLimit = endpointFileConfig.fileSizeLimit ?? mergedFileConfig.serverFileSizeLimit;
|
||||
|
||||
try {
|
||||
const formattedDate = currentDate.toISOString();
|
||||
|
|
@ -70,29 +107,143 @@ const processCodeOutput = async ({
|
|||
|
||||
const buffer = Buffer.from(response.data, 'binary');
|
||||
|
||||
const file_id = v4();
|
||||
const _file = await convertImage(req, buffer, 'high', `${file_id}${fileExt}`);
|
||||
const file = {
|
||||
..._file,
|
||||
file_id,
|
||||
usage: 1,
|
||||
// Enforce file size limit
|
||||
if (buffer.length > fileSizeLimit) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] File "${name}" (${(buffer.length / megabyte).toFixed(2)} MB) exceeds size limit of ${(fileSizeLimit / megabyte).toFixed(2)} MB, falling back to download URL`,
|
||||
);
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
|
||||
const fileIdentifier = `${session_id}/${id}`;
|
||||
|
||||
/**
|
||||
* Atomically claim a file_id for this (filename, conversationId, context) tuple.
|
||||
* Uses $setOnInsert so concurrent calls for the same filename converge on
|
||||
* a single record instead of creating duplicates (TOCTOU race fix).
|
||||
*/
|
||||
const newFileId = v4();
|
||||
const claimed = await claimCodeFile({
|
||||
filename: name,
|
||||
conversationId,
|
||||
file_id: newFileId,
|
||||
user: req.user.id,
|
||||
type: `image/${appConfig.imageOutputType}`,
|
||||
createdAt: formattedDate,
|
||||
});
|
||||
const file_id = claimed.file_id;
|
||||
const isUpdate = file_id !== newFileId;
|
||||
|
||||
if (isUpdate) {
|
||||
logger.debug(
|
||||
`[processCodeOutput] Updating existing file "${name}" (${file_id}) instead of creating duplicate`,
|
||||
);
|
||||
}
|
||||
|
||||
if (isImage) {
|
||||
const usage = isUpdate ? (claimed.usage ?? 0) + 1 : 1;
|
||||
const _file = await convertImage(req, buffer, 'high', `${file_id}${fileExt}`);
|
||||
const filepath = usage > 1 ? `${_file.filepath}?v=${Date.now()}` : _file.filepath;
|
||||
const file = {
|
||||
..._file,
|
||||
filepath,
|
||||
file_id,
|
||||
messageId,
|
||||
usage,
|
||||
filename: name,
|
||||
conversationId,
|
||||
user: req.user.id,
|
||||
type: `image/${appConfig.imageOutputType}`,
|
||||
createdAt: isUpdate ? claimed.createdAt : formattedDate,
|
||||
updatedAt: formattedDate,
|
||||
source: appConfig.fileStrategy,
|
||||
context: FileContext.execute_code,
|
||||
metadata: { fileIdentifier },
|
||||
};
|
||||
await createFile(file, true);
|
||||
return Object.assign(file, { messageId, toolCallId });
|
||||
}
|
||||
|
||||
const { saveBuffer } = getStrategyFunctions(appConfig.fileStrategy);
|
||||
if (!saveBuffer) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] saveBuffer not available for strategy ${appConfig.fileStrategy}, falling back to download URL`,
|
||||
);
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
|
||||
const detectedType = await determineFileType(buffer, true);
|
||||
const mimeType = detectedType?.mime || inferMimeType(name, '') || 'application/octet-stream';
|
||||
|
||||
/** Check MIME type support - for code-generated files, we're lenient but log unsupported types */
|
||||
const isSupportedMimeType = fileConfig.checkType(
|
||||
mimeType,
|
||||
endpointFileConfig.supportedMimeTypes,
|
||||
);
|
||||
if (!isSupportedMimeType) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] File "${name}" has unsupported MIME type "${mimeType}", proceeding with storage but may not be usable as tool resource`,
|
||||
);
|
||||
}
|
||||
|
||||
const fileName = `${file_id}__${name}`;
|
||||
const filepath = await saveBuffer({
|
||||
userId: req.user.id,
|
||||
buffer,
|
||||
fileName,
|
||||
basePath: 'uploads',
|
||||
});
|
||||
|
||||
const file = {
|
||||
file_id,
|
||||
filepath,
|
||||
messageId,
|
||||
object: 'file',
|
||||
filename: name,
|
||||
type: mimeType,
|
||||
conversationId,
|
||||
user: req.user.id,
|
||||
bytes: buffer.length,
|
||||
updatedAt: formattedDate,
|
||||
metadata: { fileIdentifier },
|
||||
source: appConfig.fileStrategy,
|
||||
context: FileContext.execute_code,
|
||||
usage: isUpdate ? (claimed.usage ?? 0) + 1 : 1,
|
||||
createdAt: isUpdate ? claimed.createdAt : formattedDate,
|
||||
};
|
||||
createFile(file, true);
|
||||
/** Note: `messageId` & `toolCallId` are not part of file DB schema; message object records associated file ID */
|
||||
|
||||
await createFile(file, true);
|
||||
return Object.assign(file, { messageId, toolCallId });
|
||||
} catch (error) {
|
||||
logAxiosError({
|
||||
message: 'Error downloading code environment file',
|
||||
message: 'Error downloading/processing code environment file',
|
||||
error,
|
||||
});
|
||||
|
||||
// Fallback for download errors - return download URL so user can still manually download
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -204,9 +355,16 @@ const primeFiles = async (options, apiKey) => {
|
|||
if (!toolContext) {
|
||||
toolContext = `- Note: The following files are available in the "${Tools.execute_code}" tool environment:`;
|
||||
}
|
||||
toolContext += `\n\t- /mnt/data/${file.filename}${
|
||||
agentResourceIds.has(file.file_id) ? '' : ' (just attached by user)'
|
||||
}`;
|
||||
|
||||
let fileSuffix = '';
|
||||
if (!agentResourceIds.has(file.file_id)) {
|
||||
fileSuffix =
|
||||
file.context === FileContext.execute_code
|
||||
? ' (from previous code execution)'
|
||||
: ' (attached by user)';
|
||||
}
|
||||
|
||||
toolContext += `\n\t- /mnt/data/${file.filename}${fileSuffix}`;
|
||||
files.push({
|
||||
id,
|
||||
session_id,
|
||||
|
|
|
|||
411
api/server/services/Files/Code/process.spec.js
Normal file
411
api/server/services/Files/Code/process.spec.js
Normal file
|
|
@ -0,0 +1,411 @@
|
|||
// Configurable file size limit for tests - use a getter so it can be changed per test
|
||||
const fileSizeLimitConfig = { value: 20 * 1024 * 1024 }; // Default 20MB
|
||||
|
||||
// Mock librechat-data-provider with configurable file size limit
|
||||
jest.mock('librechat-data-provider', () => {
|
||||
const actual = jest.requireActual('librechat-data-provider');
|
||||
return {
|
||||
...actual,
|
||||
mergeFileConfig: jest.fn((config) => {
|
||||
const merged = actual.mergeFileConfig(config);
|
||||
// Override the serverFileSizeLimit with our test value
|
||||
return {
|
||||
...merged,
|
||||
get serverFileSizeLimit() {
|
||||
return fileSizeLimitConfig.value;
|
||||
},
|
||||
};
|
||||
}),
|
||||
getEndpointFileConfig: jest.fn((options) => {
|
||||
const config = actual.getEndpointFileConfig(options);
|
||||
// Override fileSizeLimit with our test value
|
||||
return {
|
||||
...config,
|
||||
get fileSizeLimit() {
|
||||
return fileSizeLimitConfig.value;
|
||||
},
|
||||
};
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
const { FileContext } = require('librechat-data-provider');
|
||||
|
||||
// Mock uuid
|
||||
jest.mock('uuid', () => ({
|
||||
v4: jest.fn(() => 'mock-uuid-1234'),
|
||||
}));
|
||||
|
||||
// Mock axios
|
||||
jest.mock('axios');
|
||||
const axios = require('axios');
|
||||
|
||||
// Mock logger
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
warn: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock getCodeBaseURL
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
getCodeBaseURL: jest.fn(() => 'https://code-api.example.com'),
|
||||
}));
|
||||
|
||||
// Mock logAxiosError and getBasePath
|
||||
jest.mock('@librechat/api', () => ({
|
||||
logAxiosError: jest.fn(),
|
||||
getBasePath: jest.fn(() => ''),
|
||||
}));
|
||||
|
||||
// Mock models
|
||||
const mockClaimCodeFile = jest.fn();
|
||||
jest.mock('~/models', () => ({
|
||||
createFile: jest.fn().mockResolvedValue({}),
|
||||
getFiles: jest.fn(),
|
||||
updateFile: jest.fn(),
|
||||
claimCodeFile: (...args) => mockClaimCodeFile(...args),
|
||||
}));
|
||||
|
||||
// Mock permissions (must be before process.js import)
|
||||
jest.mock('~/server/services/Files/permissions', () => ({
|
||||
filterFilesByAgentAccess: jest.fn((options) => Promise.resolve(options.files)),
|
||||
}));
|
||||
|
||||
// Mock strategy functions
|
||||
jest.mock('~/server/services/Files/strategies', () => ({
|
||||
getStrategyFunctions: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock convertImage
|
||||
jest.mock('~/server/services/Files/images/convert', () => ({
|
||||
convertImage: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock determineFileType
|
||||
jest.mock('~/server/utils', () => ({
|
||||
determineFileType: jest.fn(),
|
||||
}));
|
||||
|
||||
const { createFile, getFiles } = require('~/models');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { convertImage } = require('~/server/services/Files/images/convert');
|
||||
const { determineFileType } = require('~/server/utils');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
// Import after mocks
|
||||
const { processCodeOutput } = require('./process');
|
||||
|
||||
describe('Code Process', () => {
|
||||
const mockReq = {
|
||||
user: { id: 'user-123' },
|
||||
config: {
|
||||
fileConfig: {},
|
||||
fileStrategy: 'local',
|
||||
imageOutputType: 'webp',
|
||||
},
|
||||
};
|
||||
|
||||
const baseParams = {
|
||||
req: mockReq,
|
||||
id: 'file-id-123',
|
||||
name: 'test-file.txt',
|
||||
apiKey: 'test-api-key',
|
||||
toolCallId: 'tool-call-123',
|
||||
conversationId: 'conv-123',
|
||||
messageId: 'msg-123',
|
||||
session_id: 'session-123',
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
// Default mock: atomic claim returns a new file record (no existing file)
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'mock-uuid-1234',
|
||||
user: 'user-123',
|
||||
});
|
||||
getFiles.mockResolvedValue(null);
|
||||
createFile.mockResolvedValue({});
|
||||
getStrategyFunctions.mockReturnValue({
|
||||
saveBuffer: jest.fn().mockResolvedValue('/uploads/mock-file-path.txt'),
|
||||
});
|
||||
determineFileType.mockResolvedValue({ mime: 'text/plain' });
|
||||
});
|
||||
|
||||
describe('atomic file claim (via processCodeOutput)', () => {
|
||||
it('should reuse file_id from existing record via atomic claim', async () => {
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'existing-file-id',
|
||||
filename: 'test-file.txt',
|
||||
usage: 2,
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
});
|
||||
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(mockClaimCodeFile).toHaveBeenCalledWith({
|
||||
filename: 'test-file.txt',
|
||||
conversationId: 'conv-123',
|
||||
file_id: 'mock-uuid-1234',
|
||||
user: 'user-123',
|
||||
});
|
||||
|
||||
expect(result.file_id).toBe('existing-file-id');
|
||||
expect(result.usage).toBe(3);
|
||||
expect(result.createdAt).toBe('2024-01-01T00:00:00.000Z');
|
||||
});
|
||||
|
||||
it('should create new file when no existing file found', async () => {
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'mock-uuid-1234',
|
||||
user: 'user-123',
|
||||
});
|
||||
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.file_id).toBe('mock-uuid-1234');
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('processCodeOutput', () => {
|
||||
describe('image file processing', () => {
|
||||
it('should process image files using convertImage', async () => {
|
||||
const imageParams = { ...baseParams, name: 'chart.png' };
|
||||
const imageBuffer = Buffer.alloc(500);
|
||||
axios.mockResolvedValue({ data: imageBuffer });
|
||||
|
||||
const convertedFile = {
|
||||
filepath: '/uploads/converted-image.webp',
|
||||
bytes: 400,
|
||||
};
|
||||
convertImage.mockResolvedValue(convertedFile);
|
||||
|
||||
const result = await processCodeOutput(imageParams);
|
||||
|
||||
expect(convertImage).toHaveBeenCalledWith(
|
||||
mockReq,
|
||||
imageBuffer,
|
||||
'high',
|
||||
'mock-uuid-1234.png',
|
||||
);
|
||||
expect(result.type).toBe('image/webp');
|
||||
expect(result.context).toBe(FileContext.execute_code);
|
||||
expect(result.filename).toBe('chart.png');
|
||||
});
|
||||
|
||||
it('should update existing image file with cache-busted filepath', async () => {
|
||||
const imageParams = { ...baseParams, name: 'chart.png' };
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'existing-img-id',
|
||||
usage: 1,
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
});
|
||||
|
||||
const imageBuffer = Buffer.alloc(500);
|
||||
axios.mockResolvedValue({ data: imageBuffer });
|
||||
convertImage.mockResolvedValue({ filepath: '/images/user-123/existing-img-id.webp' });
|
||||
|
||||
const result = await processCodeOutput(imageParams);
|
||||
|
||||
expect(convertImage).toHaveBeenCalledWith(
|
||||
mockReq,
|
||||
imageBuffer,
|
||||
'high',
|
||||
'existing-img-id.png',
|
||||
);
|
||||
expect(result.file_id).toBe('existing-img-id');
|
||||
expect(result.usage).toBe(2);
|
||||
expect(result.filepath).toMatch(/^\/images\/user-123\/existing-img-id\.webp\?v=\d+$/);
|
||||
expect(logger.debug).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Updating existing file'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('non-image file processing', () => {
|
||||
it('should process non-image files using saveBuffer', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const mockSaveBuffer = jest.fn().mockResolvedValue('/uploads/saved-file.txt');
|
||||
getStrategyFunctions.mockReturnValue({ saveBuffer: mockSaveBuffer });
|
||||
determineFileType.mockResolvedValue({ mime: 'text/plain' });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(mockSaveBuffer).toHaveBeenCalledWith({
|
||||
userId: 'user-123',
|
||||
buffer: smallBuffer,
|
||||
fileName: 'mock-uuid-1234__test-file.txt',
|
||||
basePath: 'uploads',
|
||||
});
|
||||
expect(result.type).toBe('text/plain');
|
||||
expect(result.filepath).toBe('/uploads/saved-file.txt');
|
||||
expect(result.bytes).toBe(100);
|
||||
});
|
||||
|
||||
it('should detect MIME type from buffer', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
determineFileType.mockResolvedValue({ mime: 'application/pdf' });
|
||||
|
||||
const result = await processCodeOutput({ ...baseParams, name: 'document.pdf' });
|
||||
|
||||
expect(determineFileType).toHaveBeenCalledWith(smallBuffer, true);
|
||||
expect(result.type).toBe('application/pdf');
|
||||
});
|
||||
|
||||
it('should fallback to application/octet-stream for unknown types', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
determineFileType.mockResolvedValue(null);
|
||||
|
||||
const result = await processCodeOutput({ ...baseParams, name: 'unknown.xyz' });
|
||||
|
||||
expect(result.type).toBe('application/octet-stream');
|
||||
});
|
||||
});
|
||||
|
||||
describe('file size limit enforcement', () => {
|
||||
it('should fallback to download URL when file exceeds size limit', async () => {
|
||||
// Set a small file size limit for this test
|
||||
fileSizeLimitConfig.value = 1000; // 1KB limit
|
||||
|
||||
const largeBuffer = Buffer.alloc(5000); // 5KB - exceeds 1KB limit
|
||||
axios.mockResolvedValue({ data: largeBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('exceeds size limit'));
|
||||
expect(result.filepath).toContain('/api/files/code/download/session-123/file-id-123');
|
||||
expect(result.expiresAt).toBeDefined();
|
||||
// Should not call createFile for oversized files (fallback path)
|
||||
expect(createFile).not.toHaveBeenCalled();
|
||||
|
||||
// Reset to default for other tests
|
||||
fileSizeLimitConfig.value = 20 * 1024 * 1024;
|
||||
});
|
||||
});
|
||||
|
||||
describe('fallback behavior', () => {
|
||||
it('should fallback to download URL when saveBuffer is not available', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
getStrategyFunctions.mockReturnValue({ saveBuffer: null });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(logger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('saveBuffer not available'),
|
||||
);
|
||||
expect(result.filepath).toContain('/api/files/code/download/');
|
||||
expect(result.filename).toBe('test-file.txt');
|
||||
});
|
||||
|
||||
it('should fallback to download URL on axios error', async () => {
|
||||
axios.mockRejectedValue(new Error('Network error'));
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.filepath).toContain('/api/files/code/download/session-123/file-id-123');
|
||||
expect(result.conversationId).toBe('conv-123');
|
||||
expect(result.messageId).toBe('msg-123');
|
||||
expect(result.toolCallId).toBe('tool-call-123');
|
||||
});
|
||||
});
|
||||
|
||||
describe('usage counter increment', () => {
|
||||
it('should set usage to 1 for new files', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
|
||||
it('should increment usage for existing files', async () => {
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'existing-id',
|
||||
usage: 5,
|
||||
createdAt: '2024-01-01',
|
||||
});
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.usage).toBe(6);
|
||||
});
|
||||
|
||||
it('should handle existing file with undefined usage', async () => {
|
||||
mockClaimCodeFile.mockResolvedValue({
|
||||
file_id: 'existing-id',
|
||||
createdAt: '2024-01-01',
|
||||
});
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('metadata and file properties', () => {
|
||||
it('should include fileIdentifier in metadata', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.metadata).toEqual({
|
||||
fileIdentifier: 'session-123/file-id-123',
|
||||
});
|
||||
});
|
||||
|
||||
it('should set correct context for code-generated files', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.context).toBe(FileContext.execute_code);
|
||||
});
|
||||
|
||||
it('should include toolCallId and messageId in result', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.toolCallId).toBe('tool-call-123');
|
||||
expect(result.messageId).toBe('msg-123');
|
||||
});
|
||||
|
||||
it('should call createFile with upsert enabled', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
await processCodeOutput(baseParams);
|
||||
|
||||
expect(createFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
file_id: 'mock-uuid-1234',
|
||||
context: FileContext.execute_code,
|
||||
}),
|
||||
true, // upsert flag
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue