From 43ff3f8473e95b1de9baecadf2f246feca029654 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 28 Feb 2026 09:06:32 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=92=B8=20fix:=20Model=20Identifier=20Edge?= =?UTF-8?q?=20Case=20in=20Agent=20Transactions=20(#11988)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🔧 fix: Add skippedAgentIds tracking in initializeClient error handling - Enhanced error handling in the initializeClient function to track agent IDs that are skipped during processing. This addition improves the ability to monitor and debug issues related to agent initialization failures. * 🔧 fix: Update model assignment in BaseClient to use instance model - Modified the model assignment in BaseClient to use `this.model` instead of `responseMessage.model`, clarifying that when using agents, the model refers to the agent ID rather than the model itself. This change improves code clarity and correctness in the context of agent usage. * 🔧 test: Add tests for recordTokenUsage model assignment in BaseClient - Introduced new test cases in BaseClient to ensure that the correct model is passed to the recordTokenUsage method, verifying that it uses this.model instead of the agent ID from responseMessage.model. This enhances the accuracy of token usage tracking in agent scenarios. - Improved error handling in the initializeClient function to log errors when processing agents, ensuring that skipped agent IDs are tracked for better debugging. --- api/app/clients/BaseClient.js | 3 +- api/app/clients/specs/BaseClient.test.js | 50 +++++++++++++++++++ .../services/Endpoints/agents/initialize.js | 8 ++- 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 2d008b9991..8f931f8a5e 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -797,7 +797,8 @@ class BaseClient { promptTokens, completionTokens, balance: balanceConfig, - model: responseMessage.model, + /** Note: When using agents, responseMessage.model is the agent ID, not the model */ + model: this.model, messageId: this.responseMessageId, }); } diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 15328af644..f13c9979ac 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -821,6 +821,56 @@ describe('BaseClient', () => { }); }); + describe('recordTokenUsage model assignment', () => { + test('should pass this.model to recordTokenUsage, not the agent ID from responseMessage.model', async () => { + const actualModel = 'claude-opus-4-5'; + const agentId = 'agent_p5Z_IU6EIxBoqn1BoqLBp'; + + TestClient.model = actualModel; + TestClient.options.endpoint = 'agents'; + TestClient.options.agent = { id: agentId }; + + TestClient.getTokenCountForResponse = jest.fn().mockReturnValue(50); + TestClient.recordTokenUsage = jest.fn().mockResolvedValue(undefined); + TestClient.buildMessages.mockReturnValue({ + prompt: [], + tokenCountMap: { res: 50 }, + }); + + await TestClient.sendMessage('Hello', {}); + + expect(TestClient.recordTokenUsage).toHaveBeenCalledWith( + expect.objectContaining({ + model: actualModel, + }), + ); + + const callArgs = TestClient.recordTokenUsage.mock.calls[0][0]; + expect(callArgs.model).not.toBe(agentId); + }); + + test('should pass this.model even when this.model differs from modelOptions.model', async () => { + const instanceModel = 'gpt-4o'; + TestClient.model = instanceModel; + TestClient.modelOptions = { model: 'gpt-4o-mini' }; + + TestClient.getTokenCountForResponse = jest.fn().mockReturnValue(50); + TestClient.recordTokenUsage = jest.fn().mockResolvedValue(undefined); + TestClient.buildMessages.mockReturnValue({ + prompt: [], + tokenCountMap: { res: 50 }, + }); + + await TestClient.sendMessage('Hello', {}); + + expect(TestClient.recordTokenUsage).toHaveBeenCalledWith( + expect.objectContaining({ + model: instanceModel, + }), + ); + }); + }); + describe('getMessagesWithinTokenLimit with instructions', () => { test('should always include instructions when present', async () => { TestClient.maxContextTokens = 50; diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index fd2e42511d..e71270ef85 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -306,6 +306,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { } } catch (err) { logger.error(`[initializeClient] Error processing agent ${agentId}:`, err); + skippedAgentIds.add(agentId); } } @@ -315,7 +316,12 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { if (checkAgentInit(agentId)) { continue; } - await processAgent(agentId); + try { + await processAgent(agentId); + } catch (err) { + logger.error(`[initializeClient] Error processing chain agent ${agentId}:`, err); + skippedAgentIds.add(agentId); + } } const chain = await createSequentialChainEdges([primaryConfig.id].concat(agent_ids), '{convo}'); collectEdges(chain);