💸 fix: Model Identifier Edge Case in Agent Transactions (#11988)

* 🔧 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.
This commit is contained in:
Danny Avila 2026-02-28 09:06:32 -05:00 committed by GitHub
parent 8b159079f5
commit 43ff3f8473
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 59 additions and 2 deletions

View file

@ -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,
});
}

View file

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

View file

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