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