🤖 refactor: Improve Agents Memory Usage, Bump Keyv, Grok 3 (#6850)

* chore: remove unused redis file

* chore: bump keyv dependencies, and update related imports

* refactor: Implement IoRedis client for rate limiting across middleware, as node-redis via keyv not compatible

* fix: Set max listeners to expected amount

* WIP: memory improvements

* refactor: Simplify getAbortData assignment in createAbortController

* refactor: Update getAbortData to use WeakRef for content management

* WIP: memory improvements in agent chat requests

* refactor: Enhance memory management with finalization registry and cleanup functions

* refactor: Simplify domainParser calls by removing unnecessary request parameter

* refactor: Update parameter types for action tools and agent loading functions to use minimal configs

* refactor: Simplify domainParser tests by removing unnecessary request parameter

* refactor: Simplify domainParser call by removing unnecessary request parameter

* refactor: Enhance client disposal by nullifying additional properties to improve memory management

* refactor: Improve title generation by adding abort controller and timeout handling, consolidate request cleanup

* refactor: Update checkIdleConnections to skip current user when checking for idle connections if passed

* refactor: Update createMCPTool to derive userId from config and handle abort signals

* refactor: Introduce createTokenCounter function and update tokenCounter usage; enhance disposeClient to reset Graph values

* refactor: Update getMCPManager to accept userId parameter for improved idle connection handling

* refactor: Extract logToolError function for improved error handling in AgentClient

* refactor: Update disposeClient to clear handlerRegistry and graphRunnable references in client.run

* refactor: Extract createHandleNewToken function to streamline token handling in initializeClient

* chore: bump @librechat/agents

* refactor: Improve timeout handling in addTitle function for better error management

* refactor: Introduce createFetch instead of using class method

* refactor: Enhance client disposal and request data handling in AskController and EditController

* refactor: Update import statements for AnthropicClient and OpenAIClient to use specific paths

* refactor: Use WeakRef for response handling in SplitStreamHandler to prevent memory leaks

* refactor: Simplify client disposal and rename getReqData to processReqData in AskController and EditController

* refactor: Improve logging structure and parameter handling in OpenAIClient

* refactor: Remove unused GraphEvents and improve stream event handling in AnthropicClient and OpenAIClient

* refactor: Simplify client initialization in AskController and EditController

* refactor: Remove unused mock functions and implement in-memory store for KeyvMongo

* chore: Update dependencies in package-lock.json to latest versions

* refactor: Await token usage recording in OpenAIClient to ensure proper async handling

* refactor: Remove handleAbort route from multiple endpoints and enhance client disposal logic

* refactor: Enhance abort controller logic by managing abortKey more effectively

* refactor: Add newConversation handling in useEventHandlers for improved conversation management

* fix: dropparams

* refactor: Use optional chaining for safer access to request properties in BaseClient

* refactor: Move client disposal and request data processing logic to cleanup module for better organization

* refactor: Remove aborted request check from addTitle function for cleaner logic

* feat: Add Grok 3 model pricing and update tests for new models

* chore: Remove trace warnings and inspect flags from backend start script used for debugging

* refactor: Replace user identifier handling with userId for consistency across controllers, use UserId in clientRegistry

* refactor: Enhance client disposal logic to prevent memory leaks by clearing additional references

* chore: Update @librechat/agents to version 2.4.14 in package.json and package-lock.json
This commit is contained in:
Danny Avila 2025-04-12 18:46:36 -04:00 committed by GitHub
parent 1e6b1b9554
commit 37964975c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
68 changed files with 1796 additions and 623 deletions

240
api/server/cleanup.js Normal file
View file

@ -0,0 +1,240 @@
const { logger } = require('~/config');
// WeakMap to hold temporary data associated with requests
const requestDataMap = new WeakMap();
const FinalizationRegistry = global.FinalizationRegistry || null;
/**
* FinalizationRegistry to clean up client objects when they are garbage collected.
* This is used to prevent memory leaks and ensure that client objects are
* properly disposed of when they are no longer needed.
* The registry holds a weak reference to the client object and a cleanup
* callback that is called when the client object is garbage collected.
* The callback can be used to perform any necessary cleanup operations,
* such as removing event listeners or freeing up resources.
*/
const clientRegistry = FinalizationRegistry
? new FinalizationRegistry((heldValue) => {
try {
// This will run when the client is garbage collected
if (heldValue && heldValue.userId) {
logger.debug(`[FinalizationRegistry] Cleaning up client for user ${heldValue.userId}`);
} else {
logger.debug('[FinalizationRegistry] Cleaning up client');
}
} catch (e) {
// Ignore errors
}
})
: null;
/**
* Cleans up the client object by removing references to its properties.
* This is useful for preventing memory leaks and ensuring that the client
* and its properties can be garbage collected when it is no longer needed.
*/
function disposeClient(client) {
if (!client) {
return;
}
try {
if (client.user) {
client.user = null;
}
if (client.apiKey) {
client.apiKey = null;
}
if (client.azure) {
client.azure = null;
}
if (client.conversationId) {
client.conversationId = null;
}
if (client.responseMessageId) {
client.responseMessageId = null;
}
if (client.clientName) {
client.clientName = null;
}
if (client.sender) {
client.sender = null;
}
if (client.model) {
client.model = null;
}
if (client.maxContextTokens) {
client.maxContextTokens = null;
}
if (client.contextStrategy) {
client.contextStrategy = null;
}
if (client.currentDateString) {
client.currentDateString = null;
}
if (client.inputTokensKey) {
client.inputTokensKey = null;
}
if (client.outputTokensKey) {
client.outputTokensKey = null;
}
if (client.run) {
// Break circular references in run
if (client.run.Graph) {
client.run.Graph.resetValues();
client.run.Graph.handlerRegistry = null;
client.run.Graph.runId = null;
client.run.Graph.tools = null;
client.run.Graph.signal = null;
client.run.Graph.config = null;
client.run.Graph.toolEnd = null;
client.run.Graph.toolMap = null;
client.run.Graph.provider = null;
client.run.Graph.streamBuffer = null;
client.run.Graph.clientOptions = null;
client.run.Graph.graphState = null;
client.run.Graph.boundModel = null;
client.run.Graph.systemMessage = null;
client.run.Graph.reasoningKey = null;
client.run.Graph.messages = null;
client.run.Graph.contentData = null;
client.run.Graph.stepKeyIds = null;
client.run.Graph.contentIndexMap = null;
client.run.Graph.toolCallStepIds = null;
client.run.Graph.messageIdsByStepKey = null;
client.run.Graph.messageStepHasToolCalls = null;
client.run.Graph.prelimMessageIdsByStepKey = null;
client.run.Graph.currentTokenType = null;
client.run.Graph.lastToken = null;
client.run.Graph.tokenTypeSwitch = null;
client.run.Graph.indexTokenCountMap = null;
client.run.Graph.currentUsage = null;
client.run.Graph.tokenCounter = null;
client.run.Graph.maxContextTokens = null;
client.run.Graph.pruneMessages = null;
client.run.Graph.lastStreamCall = null;
client.run.Graph.startIndex = null;
client.run.Graph = null;
}
if (client.run.handlerRegistry) {
client.run.handlerRegistry = null;
}
if (client.run.graphRunnable) {
if (client.run.graphRunnable.channels) {
client.run.graphRunnable.channels = null;
}
if (client.run.graphRunnable.nodes) {
client.run.graphRunnable.nodes = null;
}
if (client.run.graphRunnable.lc_kwargs) {
client.run.graphRunnable.lc_kwargs = null;
}
if (client.run.graphRunnable.builder?.nodes) {
client.run.graphRunnable.builder.nodes = null;
client.run.graphRunnable.builder = null;
}
client.run.graphRunnable = null;
}
client.run = null;
}
if (client.sendMessage) {
client.sendMessage = null;
}
if (client.savedMessageIds) {
client.savedMessageIds.clear();
client.savedMessageIds = null;
}
if (client.currentMessages) {
client.currentMessages = null;
}
if (client.streamHandler) {
client.streamHandler = null;
}
if (client.contentParts) {
client.contentParts = null;
}
if (client.abortController) {
client.abortController = null;
}
if (client.collectedUsage) {
client.collectedUsage = null;
}
if (client.indexTokenCountMap) {
client.indexTokenCountMap = null;
}
if (client.agentConfigs) {
client.agentConfigs = null;
}
if (client.artifactPromises) {
client.artifactPromises = null;
}
if (client.usage) {
client.usage = null;
}
if (typeof client.dispose === 'function') {
client.dispose();
}
if (client.options) {
if (client.options.req) {
client.options.req = null;
}
if (client.options.res) {
client.options.res = null;
}
if (client.options.attachments) {
client.options.attachments = null;
}
if (client.options.agent) {
client.options.agent = null;
}
}
client.options = null;
} catch (e) {
// Ignore errors during disposal
}
}
function processReqData(data = {}, context) {
let {
abortKey,
userMessage,
userMessagePromise,
responseMessageId,
promptTokens,
conversationId,
userMessageId,
} = context;
for (const key in data) {
if (key === 'userMessage') {
userMessage = data[key];
userMessageId = data[key].messageId;
} else if (key === 'userMessagePromise') {
userMessagePromise = data[key];
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
} else if (key === 'abortKey') {
abortKey = data[key];
} else if (!conversationId && key === 'conversationId') {
conversationId = data[key];
}
}
return {
abortKey,
userMessage,
userMessagePromise,
responseMessageId,
promptTokens,
conversationId,
userMessageId,
};
}
module.exports = {
disposeClient,
requestDataMap,
clientRegistry,
processReqData,
};