From 09c309bc78bf85ce0400a3548e3e521c4442d375 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 10 Nov 2025 13:50:17 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9E=20fix:=20Model=20End=20Callback=20?= =?UTF-8?q?and=20Streamline=20Client=20Cleanup=20(#10438)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: update agent context handling in ModelEndHandler due to new MultiAgentGraph * refactor: streamline client cleanup process by utilizing property arrays for potential circular reference removal --- api/server/cleanup.js | 139 +++++++++++++-------- api/server/controllers/agents/callbacks.js | 15 ++- 2 files changed, 97 insertions(+), 57 deletions(-) diff --git a/api/server/cleanup.js b/api/server/cleanup.js index 5a027e7f4f..f2c2b20e84 100644 --- a/api/server/cleanup.js +++ b/api/server/cleanup.js @@ -29,8 +29,59 @@ const clientRegistry = FinalizationRegistry }) : null; +const graphPropsToClean = [ + 'handlerRegistry', + 'runId', + 'tools', + 'signal', + 'config', + 'agentContexts', + 'messages', + 'contentData', + 'stepKeyIds', + 'contentIndexMap', + 'toolCallStepIds', + 'messageIdsByStepKey', + 'messageStepHasToolCalls', + 'prelimMessageIdsByStepKey', + 'startIndex', + 'defaultAgentId', + 'dispatchReasoningDelta', + 'compileOptions', + 'invokedToolIds', + 'overrideModel', +]; + +const graphRunnablePropsToClean = [ + 'lc_serializable', + 'lc_kwargs', + 'lc_runnable', + 'name', + 'lc_namespace', + 'lg_is_pregel', + 'nodes', + 'channels', + 'inputChannels', + 'outputChannels', + 'autoValidate', + 'streamMode', + 'streamChannels', + 'interruptAfter', + 'interruptBefore', + 'stepTimeout', + 'debug', + 'checkpointer', + 'retryPolicy', + 'config', + 'store', + 'triggerToNodes', + 'cache', + 'description', + 'metaRegistry', +]; + /** - * Cleans up the client object by removing references to its properties. + * Cleans up the client object by removing potential circular 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. */ @@ -223,68 +274,54 @@ function disposeClient(client) { if (client.processMemory) { client.processMemory = 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; - if (client.run.Graph.boundModel?.client) { - client.run.Graph.boundModel.client = 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; + + graphPropsToClean.forEach((prop) => { + if (client.run.Graph[prop] !== undefined) { + client.run.Graph[prop] = 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; + graphRunnablePropsToClean.forEach((prop) => { + if (client.run.graphRunnable[prop] !== undefined) { + client.run.graphRunnable[prop] = null; + } + }); + + if (client.run.graphRunnable.builder) { + if (client.run.graphRunnable.builder.nodes !== undefined) { + client.run.graphRunnable.builder.nodes = null; + } client.run.graphRunnable.builder = null; } + client.run.graphRunnable = null; } + + const runPropsToClean = [ + 'handlerRegistry', + 'id', + 'indexTokenCountMap', + 'returnContent', + 'tokenCounter', + ]; + + runPropsToClean.forEach((prop) => { + if (client.run[prop] !== undefined) { + client.run[prop] = null; + } + }); + client.run = null; } + if (client.sendMessage) { client.sendMessage = null; } diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index a66fe9a053..411c4ed9c5 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -41,7 +41,11 @@ class ModelEndHandler { } try { - if (metadata.provider === Providers.GOOGLE || graph.clientOptions?.disableStreaming) { + const agentContext = graph.getAgentContext(metadata); + if ( + agentContext.provider === Providers.GOOGLE || + agentContext.clientOptions?.disableStreaming + ) { handleToolCalls(data?.output?.tool_calls, metadata, graph); } @@ -49,14 +53,13 @@ class ModelEndHandler { if (!usage) { return; } - if (metadata?.model) { - usage.model = metadata.model; + const modelName = metadata?.ls_model_name || agentContext.clientOptions?.model; + if (modelName) { + usage.model = modelName; } this.collectedUsage.push(usage); - const streamingDisabled = !!( - graph.clientOptions?.disableStreaming || graph?.boundModel?.disableStreaming - ); + const streamingDisabled = !!agentContext.clientOptions?.disableStreaming; if (!streamingDisabled) { return; }