mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-27 04:07:05 +00:00
🧾 fix: Report Complete Agents API Usage (#15127)
* fix: report complete agents api usage * fix: preserve invoked usage context * test: cover absent usage context * fix: type responses usage finalization * fix: preserve reasoning usage aliases * fix: declare reasoning usage alias
This commit is contained in:
Danny Avila
GitHub
parent
719b04f389
commit
dd146ff74d
17 changed files with 581 additions and 98 deletions
|
|
@ -18,7 +18,7 @@ jest.mock('~/server/services/Files/process', () => ({
|
|||
saveBase64Image: jest.fn(),
|
||||
}));
|
||||
|
||||
const { ModelEndHandler } = require('../callbacks');
|
||||
const { ModelEndHandler, contextualizeModelUsage } = require('../callbacks');
|
||||
|
||||
const buildGraph = () => ({
|
||||
getAgentContext: () => ({
|
||||
|
|
@ -28,6 +28,46 @@ const buildGraph = () => ({
|
|||
});
|
||||
|
||||
describe('ModelEndHandler — Vertex thoughtSignature capture (issue #13006 follow-up)', () => {
|
||||
it('leaves usage usable when graph context is unavailable', () => {
|
||||
const usage = { input_tokens: 10, output_tokens: 5 };
|
||||
|
||||
expect(contextualizeModelUsage(usage, undefined, undefined)).toEqual(usage);
|
||||
expect(contextualizeModelUsage(usage, undefined, null)).toEqual(usage);
|
||||
});
|
||||
|
||||
it('prefers the actually invoked fallback provider and model', () => {
|
||||
const usage = { input_tokens: 10, output_tokens: 5 };
|
||||
const result = contextualizeModelUsage(
|
||||
usage,
|
||||
{
|
||||
__invoked_provider: 'anthropic',
|
||||
__invoked_model: 'claude-fallback',
|
||||
},
|
||||
{
|
||||
provider: 'bedrock',
|
||||
agentId: 'agent-1',
|
||||
clientOptions: { model: 'configured-model' },
|
||||
},
|
||||
);
|
||||
|
||||
expect(result).toEqual({
|
||||
...usage,
|
||||
provider: 'anthropic',
|
||||
model: 'claude-fallback',
|
||||
agentId: 'agent-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('prefers provider-reported model metadata over the invoked fallback model', () => {
|
||||
expect(
|
||||
contextualizeModelUsage(
|
||||
{ input_tokens: 10, output_tokens: 5 },
|
||||
{ ls_model_name: 'reported-model', __invoked_model: 'fallback-model' },
|
||||
{ clientOptions: { model: 'configured-model' } },
|
||||
).model,
|
||||
).toBe('reported-model');
|
||||
});
|
||||
|
||||
it('maps non-empty signatures onto tool_call_ids in order', async () => {
|
||||
const collectedUsage = [];
|
||||
const collectedThoughtSignatures = {};
|
||||
|
|
@ -170,6 +210,8 @@ describe('ModelEndHandler — Vertex thoughtSignature capture (issue #13006 foll
|
|||
);
|
||||
|
||||
expect(collectedUsage[0].agentId).toBe('agent_sub');
|
||||
expect(collectedUsage[0].provider).toBe('openai');
|
||||
expect(collectedUsage[0].model).toBe('gpt-4');
|
||||
expect(emitUsage).toHaveBeenCalledWith(expect.objectContaining({ agentId: 'agent_sub' }));
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,14 @@ const mockBuildAgentScopedContext = jest.fn().mockResolvedValue(new Map());
|
|||
const mockBuildAgentContextAttachmentsByAgentId = jest.fn().mockReturnValue(new Map());
|
||||
const mockBuildInlineMemoryContext = jest.fn().mockResolvedValue('');
|
||||
const mockApplyContextToAgent = jest.fn().mockResolvedValue(undefined);
|
||||
const mockCompletionUsage = {
|
||||
prompt_tokens: 125,
|
||||
completion_tokens: 50,
|
||||
total_tokens: 175,
|
||||
primary: { prompt_tokens: 100, completion_tokens: 40, total_tokens: 140 },
|
||||
subagent: { prompt_tokens: 25, completion_tokens: 10, total_tokens: 35 },
|
||||
};
|
||||
const mockBuildCompletionUsage = jest.fn().mockReturnValue(mockCompletionUsage);
|
||||
const mockInitialSessions = new Map([['execute_code', { session_id: 'seeded' }]]);
|
||||
const mockGetSafeErrorMetadata = jest.fn((error) => {
|
||||
const status = error?.status ?? error?.statusCode ?? error?.response?.status;
|
||||
|
|
@ -190,6 +198,7 @@ jest.mock('@librechat/api', () => ({
|
|||
.mockImplementation(({ accessibleSkillIds }) => accessibleSkillIds),
|
||||
loadSkillStates: jest.fn().mockResolvedValue({ skillStates: {}, defaultActiveOnShare: false }),
|
||||
sendFinalChunk: jest.fn(),
|
||||
buildCompletionUsage: mockBuildCompletionUsage,
|
||||
createSafeUser: jest.fn().mockReturnValue({ id: 'user-123' }),
|
||||
validateRequest: jest
|
||||
.fn()
|
||||
|
|
@ -312,7 +321,7 @@ const mockGetCacheMultiplier = jest.fn().mockReturnValue(null);
|
|||
jest.mock('~/server/controllers/agents/callbacks', () => ({
|
||||
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
buildSummarizationHandlers: jest.fn().mockReturnValue({}),
|
||||
markSummarizationUsage: jest.fn().mockImplementation((usage) => usage),
|
||||
contextualizeModelUsage: jest.fn().mockImplementation((usage) => usage),
|
||||
agentLogHandlerObj: { handle: jest.fn() },
|
||||
}));
|
||||
|
||||
|
|
@ -462,15 +471,34 @@ describe('OpenAIChatCompletionController', () => {
|
|||
expect(createRun).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ initialSessions: mockInitialSessions }),
|
||||
);
|
||||
expect(createSubagentUsageSink).toHaveBeenCalledWith(expect.any(Array), expect.any(Function));
|
||||
const aggregator =
|
||||
require('@librechat/api').createOpenAIContentAggregator.mock.results.at(-1).value;
|
||||
const initialPromptTokens = aggregator.usage.promptTokens;
|
||||
const initialCompletionTokens = aggregator.usage.completionTokens;
|
||||
const onSubagentUsage = createSubagentUsageSink.mock.calls.at(-1)[1];
|
||||
onSubagentUsage({ input_tokens: 25, output_tokens: 10 });
|
||||
expect(aggregator.usage.promptTokens).toBe(initialPromptTokens + 25);
|
||||
expect(aggregator.usage.completionTokens).toBe(initialCompletionTokens + 10);
|
||||
expect(createSubagentUsageSink).toHaveBeenCalledWith(expect.any(Array));
|
||||
});
|
||||
|
||||
it('uses collected usage for the non-streaming response', async () => {
|
||||
const { buildNonStreamingResponse } = require('@librechat/api');
|
||||
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
const collectedUsage = mockRecordCollectedUsage.mock.calls.at(-1)[1].collectedUsage;
|
||||
expect(mockBuildCompletionUsage).toHaveBeenCalledWith(collectedUsage);
|
||||
expect(buildNonStreamingResponse).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
expect.anything(),
|
||||
expect.anything(),
|
||||
expect.anything(),
|
||||
mockCompletionUsage,
|
||||
);
|
||||
});
|
||||
|
||||
it('uses collected usage in the final streaming chunk', async () => {
|
||||
const { validateRequest, sendFinalChunk } = require('@librechat/api');
|
||||
validateRequest.mockReturnValueOnce({
|
||||
request: { model: 'agent-123', messages: [], stream: true },
|
||||
});
|
||||
|
||||
await OpenAIChatCompletionController(req, res);
|
||||
|
||||
expect(sendFinalChunk).toHaveBeenCalledWith(expect.anything(), 'stop', mockCompletionUsage);
|
||||
});
|
||||
|
||||
describe('content filtering', () => {
|
||||
|
|
|
|||
|
|
@ -138,6 +138,16 @@ const mockBuildAgentScopedContext = jest.fn().mockResolvedValue(new Map());
|
|||
const mockBuildAgentContextAttachmentsByAgentId = jest.fn().mockReturnValue(new Map());
|
||||
const mockBuildInlineMemoryContext = jest.fn().mockResolvedValue('');
|
||||
const mockApplyContextToAgent = jest.fn().mockResolvedValue(undefined);
|
||||
const mockResponsesUsage = {
|
||||
input_tokens: 125,
|
||||
output_tokens: 50,
|
||||
total_tokens: 175,
|
||||
input_tokens_details: { cached_tokens: 0 },
|
||||
output_tokens_details: { reasoning_tokens: 0 },
|
||||
primary: { input_tokens: 100, output_tokens: 40, total_tokens: 140 },
|
||||
subagent: { input_tokens: 25, output_tokens: 10, total_tokens: 35 },
|
||||
};
|
||||
const mockBuildResponsesUsage = jest.fn().mockReturnValue(mockResponsesUsage);
|
||||
|
||||
jest.mock('nanoid', () => ({
|
||||
nanoid: jest.fn(() => 'mock-nanoid-123'),
|
||||
|
|
@ -294,7 +304,7 @@ jest.mock('@librechat/api', () => ({
|
|||
emitResponseCreated: jest.fn(),
|
||||
createResponseContext: jest.fn().mockReturnValue({ responseId: 'resp_123' }),
|
||||
createResponseTracker: jest.fn().mockReturnValue({
|
||||
usage: { promptTokens: 100, completionTokens: 50 },
|
||||
usage: { inputTokens: 100, outputTokens: 50, reasoningTokens: 0, cachedTokens: 0 },
|
||||
}),
|
||||
setupStreamingResponse: jest.fn(),
|
||||
emitResponseInProgress: jest.fn(),
|
||||
|
|
@ -308,8 +318,9 @@ jest.mock('@librechat/api', () => ({
|
|||
output: [],
|
||||
usage: { input_tokens: 100, output_tokens: 50, total_tokens: 150 },
|
||||
}),
|
||||
buildResponsesUsage: mockBuildResponsesUsage,
|
||||
createResponseAggregator: jest.fn().mockReturnValue({
|
||||
usage: { promptTokens: 100, completionTokens: 50 },
|
||||
usage: { inputTokens: 100, outputTokens: 50, reasoningTokens: 0, cachedTokens: 0 },
|
||||
}),
|
||||
sendResponsesErrorResponse: jest.fn(),
|
||||
createResponsesEventHandlers: jest.fn().mockReturnValue({
|
||||
|
|
@ -346,7 +357,7 @@ jest.mock('~/server/controllers/agents/callbacks', () => {
|
|||
return {
|
||||
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
createResponsesToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
|
||||
markSummarizationUsage: jest.fn().mockImplementation((usage) => usage),
|
||||
contextualizeModelUsage: jest.fn().mockImplementation((usage) => usage),
|
||||
agentLogHandlerObj: noop,
|
||||
buildSummarizationHandlers: jest.fn().mockReturnValue({
|
||||
on_summarize_start: noop,
|
||||
|
|
@ -617,6 +628,7 @@ describe('createResponse controller', () => {
|
|||
isCreatedByUser: false,
|
||||
langfuseSampled: true,
|
||||
langfuseDestinationIds: ['destination-1'],
|
||||
tokenCount: 50,
|
||||
}),
|
||||
{ context: 'Responses API - save assistant response' },
|
||||
);
|
||||
|
|
@ -1880,24 +1892,35 @@ describe('createResponse controller', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('adds subagent usage to the response usage handler', async () => {
|
||||
it('uses collected usage for the non-streaming response', async () => {
|
||||
const api = require('@librechat/api');
|
||||
api.validateResponseRequest.mockReturnValueOnce({
|
||||
request: { model: 'agent-123', input: 'Hello', stream: false },
|
||||
});
|
||||
|
||||
await createResponse(req, res);
|
||||
|
||||
const onSubagentUsage = api.createSubagentUsageSink.mock.calls.at(-1)[1];
|
||||
const aggregatorHandlers =
|
||||
api.createAggregatorEventHandlers.mock.results.at(-1)?.value ??
|
||||
api.createResponsesEventHandlers.mock.results.at(-1)?.value.handlers;
|
||||
onSubagentUsage({ input_tokens: 25, output_tokens: 10 });
|
||||
|
||||
expect(aggregatorHandlers.on_chat_model_end.handle).toHaveBeenCalledWith(
|
||||
'on_chat_model_end',
|
||||
{
|
||||
output: { usage_metadata: { input_tokens: 25, output_tokens: 10 } },
|
||||
},
|
||||
const collectedUsage = mockRecordCollectedUsage.mock.calls.at(-1)[1].collectedUsage;
|
||||
expect(mockBuildResponsesUsage).toHaveBeenCalledWith(collectedUsage);
|
||||
expect(api.buildAggregatedResponse).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
expect.anything(),
|
||||
mockResponsesUsage,
|
||||
);
|
||||
});
|
||||
|
||||
it('uses collected usage for the completed streaming event', async () => {
|
||||
const api = require('@librechat/api');
|
||||
api.validateResponseRequest.mockReturnValueOnce({
|
||||
request: { model: 'agent-123', input: 'Hello', stream: true },
|
||||
});
|
||||
|
||||
await createResponse(req, res);
|
||||
|
||||
const finalizeStream =
|
||||
api.createResponsesEventHandlers.mock.results.at(-1).value.finalizeStream;
|
||||
expect(finalizeStream).toHaveBeenCalledWith(mockResponsesUsage);
|
||||
});
|
||||
});
|
||||
|
||||
describe('sub-agent skill priming', () => {
|
||||
|
|
|
|||
|
|
@ -149,20 +149,7 @@ class ModelEndHandler {
|
|||
if (!usage) {
|
||||
return this.finalize(errorMessage);
|
||||
}
|
||||
const modelName = metadata?.ls_model_name || agentContext.clientOptions?.model;
|
||||
if (modelName) {
|
||||
usage.model = modelName;
|
||||
}
|
||||
if (agentContext.provider) {
|
||||
usage.provider = agentContext.provider;
|
||||
}
|
||||
/** Tag the producing agent so multi-endpoint graphs can price each call
|
||||
* with its own endpoint token config (recordCollectedUsage resolver). */
|
||||
if (agentContext.agentId) {
|
||||
usage.agentId = agentContext.agentId;
|
||||
}
|
||||
|
||||
let taggedUsage = markSummarizationUsage(usage, metadata);
|
||||
let taggedUsage = contextualizeModelUsage(usage, metadata, agentContext);
|
||||
/** Hidden intermediate sequential-agent calls are billed but never shown.
|
||||
* Tag them non-primary on the COLLECTED usage too (not just the emit) so
|
||||
* recordCollectedUsage excludes their output from the parent's tokenCount
|
||||
|
|
@ -1442,6 +1429,35 @@ function markSummarizationUsage(usage, metadata) {
|
|||
return usage;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stamps provider/model/agent identity onto one model call before billing or
|
||||
* API response aggregation. The graph owns this context; provider payloads do
|
||||
* not consistently include it, and cache normalization depends on it.
|
||||
*/
|
||||
function contextualizeModelUsage(usage, metadata, agentContext = {}) {
|
||||
const taggedUsage = { ...usage };
|
||||
const context = agentContext ?? {};
|
||||
const invokedProvider = metadata?.__invoked_provider;
|
||||
const invokedModel = metadata?.__invoked_model;
|
||||
const modelName =
|
||||
metadata?.ls_model_name ||
|
||||
(typeof invokedModel === 'string' && invokedModel !== '' ? invokedModel : undefined) ||
|
||||
context.clientOptions?.model;
|
||||
const provider =
|
||||
(typeof invokedProvider === 'string' && invokedProvider !== '' ? invokedProvider : undefined) ||
|
||||
context.provider;
|
||||
if (modelName) {
|
||||
taggedUsage.model = modelName;
|
||||
}
|
||||
if (provider) {
|
||||
taggedUsage.provider = provider;
|
||||
}
|
||||
if (context.agentId) {
|
||||
taggedUsage.agentId = context.agentId;
|
||||
}
|
||||
return markSummarizationUsage(taggedUsage, metadata);
|
||||
}
|
||||
|
||||
const agentLogHandlerObj = { handle: agentLogHandler };
|
||||
|
||||
/**
|
||||
|
|
@ -1480,6 +1496,7 @@ module.exports = {
|
|||
createBackgroundCodeResultHandler,
|
||||
isStreamWritable,
|
||||
markSummarizationUsage,
|
||||
contextualizeModelUsage,
|
||||
buildSummarizationHandlers,
|
||||
createResponsesToolEndCallback,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ const {
|
|||
createMCPRuntimeRequestBody,
|
||||
loadSkillStates,
|
||||
sendFinalChunk,
|
||||
buildCompletionUsage,
|
||||
createSafeUser,
|
||||
validateRequest,
|
||||
initializeAgent,
|
||||
|
|
@ -62,7 +63,7 @@ const {
|
|||
} = require('@librechat/api');
|
||||
const {
|
||||
buildSummarizationHandlers,
|
||||
markSummarizationUsage,
|
||||
contextualizeModelUsage,
|
||||
createToolEndCallback,
|
||||
agentLogHandlerObj,
|
||||
} = require('~/server/controllers/agents/callbacks');
|
||||
|
|
@ -659,12 +660,6 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
|
|||
// Create tracker for streaming or aggregator for non-streaming
|
||||
const tracker = isStreaming ? createOpenAIStreamTracker() : null;
|
||||
const aggregator = isStreaming ? null : createOpenAIContentAggregator();
|
||||
const accumulateResponseUsage = (usage) => {
|
||||
const target = isStreaming ? tracker : aggregator;
|
||||
target.usage.promptTokens += usage.input_tokens ?? 0;
|
||||
target.usage.completionTokens += usage.output_tokens ?? 0;
|
||||
};
|
||||
|
||||
// Set up response for streaming
|
||||
if (isStreaming) {
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
|
|
@ -915,12 +910,12 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
|
|||
|
||||
// Usage tracking
|
||||
on_chat_model_end: {
|
||||
handle: (_event, data, metadata) => {
|
||||
handle: (_event, data, metadata, graph) => {
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
const taggedUsage = markSummarizationUsage(usage, metadata);
|
||||
const agentContext = graph?.getAgentContext?.(metadata);
|
||||
const taggedUsage = contextualizeModelUsage(usage, metadata, agentContext);
|
||||
collectedUsage.push(taggedUsage);
|
||||
accumulateResponseUsage(taggedUsage);
|
||||
}
|
||||
},
|
||||
},
|
||||
|
|
@ -1000,7 +995,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
|
|||
tenantId: principal.tenantId,
|
||||
/** Bills subagent child-run model calls (reported outside the
|
||||
* streamEvents loop) into the same collectedUsage array. */
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage, accumulateResponseUsage),
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage),
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
|
|
@ -1056,10 +1051,12 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
|
|||
logger.error('[OpenAI API] Error recording usage:', getSafeErrorMetadata(err));
|
||||
});
|
||||
|
||||
const usage = buildCompletionUsage(collectedUsage);
|
||||
|
||||
// Finalize response
|
||||
const duration = Date.now() - requestStartTime;
|
||||
if (isStreaming) {
|
||||
sendFinalChunk(handlerConfig);
|
||||
sendFinalChunk(handlerConfig, 'stop', usage);
|
||||
res.end();
|
||||
logger.debug(`[OpenAI API] Response ${responseId} completed in ${duration}ms (streaming)`);
|
||||
|
||||
|
|
@ -1085,19 +1082,6 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
|
|||
}
|
||||
}
|
||||
|
||||
// Build usage from aggregated data
|
||||
const usage = {
|
||||
prompt_tokens: aggregator.usage.promptTokens,
|
||||
completion_tokens: aggregator.usage.completionTokens,
|
||||
total_tokens: aggregator.usage.promptTokens + aggregator.usage.completionTokens,
|
||||
};
|
||||
|
||||
if (aggregator.usage.reasoningTokens > 0) {
|
||||
usage.completion_tokens_details = {
|
||||
reasoning_tokens: aggregator.usage.reasoningTokens,
|
||||
};
|
||||
}
|
||||
|
||||
const response = buildNonStreamingResponse(
|
||||
context,
|
||||
aggregator.getText(),
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ const {
|
|||
convertInputToMessages,
|
||||
validateResponseRequest,
|
||||
buildAggregatedResponse,
|
||||
buildResponsesUsage,
|
||||
createResponseAggregator,
|
||||
sendResponsesErrorResponse,
|
||||
createResponsesEventHandlers,
|
||||
|
|
@ -77,7 +78,7 @@ const {
|
|||
const {
|
||||
createResponsesToolEndCallback,
|
||||
buildSummarizationHandlers,
|
||||
markSummarizationUsage,
|
||||
contextualizeModelUsage,
|
||||
createToolEndCallback,
|
||||
agentLogHandlerObj,
|
||||
} = require('~/server/controllers/agents/callbacks');
|
||||
|
|
@ -355,9 +356,17 @@ async function saveInputMessages(req, conversationId, inputMessages, agentId) {
|
|||
* @param {string} responseId
|
||||
* @param {import('@librechat/api').Response} response
|
||||
* @param {string} agentId
|
||||
* @param {number | undefined} visibleOutputTokens
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function saveResponseOutput(req, conversationId, responseId, response, agentId) {
|
||||
async function saveResponseOutput(
|
||||
req,
|
||||
conversationId,
|
||||
responseId,
|
||||
response,
|
||||
agentId,
|
||||
visibleOutputTokens,
|
||||
) {
|
||||
// Extract text content from output items
|
||||
let responseText = '';
|
||||
for (const item of response.output) {
|
||||
|
|
@ -386,7 +395,7 @@ async function saveResponseOutput(req, conversationId, responseId, response, age
|
|||
endpoint: EModelEndpoint.agents,
|
||||
model: agentId,
|
||||
finish_reason: response.status === 'completed' ? 'stop' : response.status,
|
||||
tokenCount: response.usage?.output_tokens,
|
||||
tokenCount: visibleOutputTokens ?? response.usage?.output_tokens,
|
||||
},
|
||||
{ context: 'Responses API - save assistant response' },
|
||||
);
|
||||
|
|
@ -1056,11 +1065,12 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
on_run_step: responsesHandlers.on_run_step,
|
||||
on_run_step_delta: responsesHandlers.on_run_step_delta,
|
||||
on_chat_model_end: {
|
||||
handle: (event, data, metadata) => {
|
||||
handle: (event, data, metadata, graph) => {
|
||||
responsesHandlers.on_chat_model_end.handle(event, data);
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
const taggedUsage = markSummarizationUsage(usage, metadata);
|
||||
const agentContext = graph?.getAgentContext?.(metadata);
|
||||
const taggedUsage = contextualizeModelUsage(usage, metadata, agentContext);
|
||||
collectedUsage.push(taggedUsage);
|
||||
}
|
||||
},
|
||||
|
|
@ -1098,11 +1108,7 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
tenantId: principal.tenantId,
|
||||
/** Bills subagent child-run model calls (reported outside the
|
||||
* streamEvents loop) into the same collectedUsage array. */
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage, (usage) => {
|
||||
responsesHandlers.on_chat_model_end.handle('on_chat_model_end', {
|
||||
output: { usage_metadata: usage },
|
||||
});
|
||||
}),
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage),
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
|
|
@ -1158,8 +1164,10 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
logger.error('[Responses API] Error recording usage:', getSafeErrorMetadata(err));
|
||||
});
|
||||
|
||||
const usage = buildResponsesUsage(collectedUsage);
|
||||
|
||||
// Finalize the stream
|
||||
finalizeStream();
|
||||
finalizeStream(usage);
|
||||
res.end();
|
||||
|
||||
const duration = Date.now() - requestStartTime;
|
||||
|
|
@ -1176,7 +1184,14 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
|
||||
// Build response for saving (use tracker with buildResponse for streaming)
|
||||
const finalResponse = buildResponse(context, tracker, 'completed');
|
||||
await saveResponseOutput(req, conversationId, responseId, finalResponse, agentId);
|
||||
await saveResponseOutput(
|
||||
req,
|
||||
conversationId,
|
||||
responseId,
|
||||
finalResponse,
|
||||
agentId,
|
||||
tracker.usage.outputTokens,
|
||||
);
|
||||
|
||||
logger.debug(
|
||||
`[Responses API] Stored response ${responseId} in conversation ${conversationId}`,
|
||||
|
|
@ -1246,11 +1261,12 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
on_run_step: aggregatorHandlers.on_run_step,
|
||||
on_run_step_delta: aggregatorHandlers.on_run_step_delta,
|
||||
on_chat_model_end: {
|
||||
handle: (event, data, metadata) => {
|
||||
handle: (event, data, metadata, graph) => {
|
||||
aggregatorHandlers.on_chat_model_end.handle(event, data);
|
||||
const usage = data?.output?.usage_metadata;
|
||||
if (usage) {
|
||||
const taggedUsage = markSummarizationUsage(usage, metadata);
|
||||
const agentContext = graph?.getAgentContext?.(metadata);
|
||||
const taggedUsage = contextualizeModelUsage(usage, metadata, agentContext);
|
||||
collectedUsage.push(taggedUsage);
|
||||
}
|
||||
},
|
||||
|
|
@ -1287,11 +1303,7 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
tenantId: principal.tenantId,
|
||||
/** Bills subagent child-run model calls (reported outside the
|
||||
* streamEvents loop) into the same collectedUsage array. */
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage, (usage) => {
|
||||
aggregatorHandlers.on_chat_model_end.handle('on_chat_model_end', {
|
||||
output: { usage_metadata: usage },
|
||||
});
|
||||
}),
|
||||
subagentUsageSink: createSubagentUsageSink(collectedUsage),
|
||||
});
|
||||
|
||||
if (!run) {
|
||||
|
|
@ -1357,7 +1369,11 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
}
|
||||
}
|
||||
|
||||
const response = buildAggregatedResponse(context, aggregator);
|
||||
const response = buildAggregatedResponse(
|
||||
context,
|
||||
aggregator,
|
||||
buildResponsesUsage(collectedUsage),
|
||||
);
|
||||
|
||||
if (request.store === true) {
|
||||
try {
|
||||
|
|
@ -1365,7 +1381,14 @@ const executeResponse = async (envelope, { req, res }) => {
|
|||
|
||||
await saveInputMessages(req, conversationId, inputMessages, agentId);
|
||||
|
||||
await saveResponseOutput(req, conversationId, responseId, response, agentId);
|
||||
await saveResponseOutput(
|
||||
req,
|
||||
conversationId,
|
||||
responseId,
|
||||
response,
|
||||
agentId,
|
||||
aggregator.usage.outputTokens,
|
||||
);
|
||||
|
||||
logger.debug(
|
||||
`[Responses API] Stored response ${responseId} in conversation ${conversationId}`,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,12 @@
|
|||
import type { Response as ServerResponse } from 'express';
|
||||
import type { UsageMetadata } from '~/stream/interfaces/IJobStore';
|
||||
import type { OpenAIResponseContext } from './types';
|
||||
import { sendFinalChunk, OpenAIModelEndHandler, createOpenAIStreamTracker } from './handlers';
|
||||
import {
|
||||
sendFinalChunk,
|
||||
buildCompletionUsage,
|
||||
OpenAIModelEndHandler,
|
||||
createOpenAIStreamTracker,
|
||||
} from './handlers';
|
||||
|
||||
describe('OpenAI-compatible agent stream handlers', () => {
|
||||
const context: OpenAIResponseContext = {
|
||||
|
|
@ -62,4 +68,56 @@ describe('OpenAI-compatible agent stream handlers', () => {
|
|||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('streams the collected primary and subagent usage override', () => {
|
||||
const tracker = createOpenAIStreamTracker();
|
||||
const writes: string[] = [];
|
||||
const res = {
|
||||
write: (chunk: string) => {
|
||||
writes.push(chunk);
|
||||
},
|
||||
} as unknown as ServerResponse;
|
||||
const usage = buildCompletionUsage([
|
||||
{ input_tokens: 100, output_tokens: 40, provider: 'openai' },
|
||||
{
|
||||
input_tokens: 25,
|
||||
output_tokens: 10,
|
||||
provider: 'openai',
|
||||
usage_type: 'subagent',
|
||||
},
|
||||
]);
|
||||
|
||||
sendFinalChunk({ context, tracker, res }, 'stop', usage);
|
||||
|
||||
const finalChunk = JSON.parse(writes[0].replace(/^data: /, '').trim());
|
||||
expect(finalChunk.usage).toEqual({
|
||||
prompt_tokens: 125,
|
||||
completion_tokens: 50,
|
||||
total_tokens: 175,
|
||||
primary: { prompt_tokens: 100, completion_tokens: 40, total_tokens: 140 },
|
||||
subagent: { prompt_tokens: 25, completion_tokens: 10, total_tokens: 35 },
|
||||
});
|
||||
});
|
||||
|
||||
it('snapshots completed response usage before later detached calls arrive', () => {
|
||||
const collectedUsage: UsageMetadata[] = [
|
||||
{ input_tokens: 100, output_tokens: 40, provider: 'openAI' },
|
||||
];
|
||||
const completedUsage = buildCompletionUsage(collectedUsage);
|
||||
|
||||
collectedUsage.push({
|
||||
input_tokens: 25,
|
||||
output_tokens: 10,
|
||||
provider: 'openAI',
|
||||
usage_type: 'subagent',
|
||||
});
|
||||
|
||||
expect(completedUsage).toEqual({
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 40,
|
||||
total_tokens: 140,
|
||||
primary: { prompt_tokens: 100, completion_tokens: 40, total_tokens: 140 },
|
||||
subagent: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -12,8 +12,10 @@ import type {
|
|||
CompletionUsage,
|
||||
ToolCall,
|
||||
} from './types';
|
||||
import type { UsageMetadata } from '~/stream/interfaces/IJobStore';
|
||||
import type { ToolExecuteOptions } from '~/agents/handlers';
|
||||
import { createToolExecuteHandler } from '~/agents/handlers';
|
||||
import { aggregateCollectedUsage } from '../usage';
|
||||
|
||||
/**
|
||||
* Create a chat completion chunk in OpenAI format
|
||||
|
|
@ -439,6 +441,7 @@ export function createOpenAIHandlers(
|
|||
export function sendFinalChunk(
|
||||
config: OpenAIStreamHandlerConfig,
|
||||
finishReason: ChatCompletionChunkChoice['finish_reason'] = 'stop',
|
||||
usageOverride?: CompletionUsage,
|
||||
): void {
|
||||
const { res, context, tracker } = config;
|
||||
|
||||
|
|
@ -449,14 +452,14 @@ export function sendFinalChunk(
|
|||
}
|
||||
|
||||
// Build usage object with reasoning token details (OpenRouter/OpenAI convention)
|
||||
const usage: CompletionUsage = {
|
||||
const usage: CompletionUsage = usageOverride ?? {
|
||||
prompt_tokens: tracker.usage.promptTokens,
|
||||
completion_tokens: tracker.usage.completionTokens,
|
||||
total_tokens: tracker.usage.promptTokens + tracker.usage.completionTokens,
|
||||
};
|
||||
|
||||
// Add reasoning token breakdown if there are reasoning tokens
|
||||
if (tracker.usage.reasoningTokens > 0) {
|
||||
if (usageOverride == null && tracker.usage.reasoningTokens > 0) {
|
||||
usage.completion_tokens_details = {
|
||||
reasoning_tokens: tracker.usage.reasoningTokens,
|
||||
};
|
||||
|
|
@ -468,3 +471,28 @@ export function sendFinalChunk(
|
|||
// Send [DONE] marker
|
||||
writeSSE(res, '[DONE]');
|
||||
}
|
||||
|
||||
/** Build provider-normalized chat-completion usage from every billed call. */
|
||||
export function buildCompletionUsage(
|
||||
collectedUsage: ReadonlyArray<UsageMetadata | null | undefined>,
|
||||
): CompletionUsage {
|
||||
const { total, primary, subagent } = aggregateCollectedUsage(collectedUsage);
|
||||
return {
|
||||
prompt_tokens: total.inputTokens,
|
||||
completion_tokens: total.outputTokens,
|
||||
total_tokens: total.totalTokens,
|
||||
...(total.reasoningTokens > 0 && {
|
||||
completion_tokens_details: { reasoning_tokens: total.reasoningTokens },
|
||||
}),
|
||||
primary: {
|
||||
prompt_tokens: primary.inputTokens,
|
||||
completion_tokens: primary.outputTokens,
|
||||
total_tokens: primary.totalTokens,
|
||||
},
|
||||
subagent: {
|
||||
prompt_tokens: subagent.inputTokens,
|
||||
completion_tokens: subagent.outputTokens,
|
||||
total_tokens: subagent.totalTokens,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -80,6 +80,12 @@ export interface ChatCompletionRequest {
|
|||
/**
|
||||
* Token usage information
|
||||
*/
|
||||
export interface CompletionUsageTotals {
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
total_tokens: number;
|
||||
}
|
||||
|
||||
export interface CompletionUsage {
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
|
|
@ -88,6 +94,10 @@ export interface CompletionUsage {
|
|||
completion_tokens_details?: {
|
||||
reasoning_tokens?: number;
|
||||
};
|
||||
/** LibreChat extension for parent, handoff, and summarization model calls. */
|
||||
primary?: CompletionUsageTotals;
|
||||
/** LibreChat extension for isolated subagent child model calls. */
|
||||
subagent?: CompletionUsageTotals;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -1,11 +1,22 @@
|
|||
import type { InputItem } from '../types';
|
||||
import type { Response as ServerResponse } from 'express';
|
||||
import type { InputItem, ResponseContext } from '../types';
|
||||
import {
|
||||
buildAggregatedResponse,
|
||||
convertInputToMessages,
|
||||
createAggregatorEventHandlers,
|
||||
createResponseAggregator,
|
||||
createResponsesEventHandlers,
|
||||
buildResponsesUsage,
|
||||
} from '../service';
|
||||
import { createResponseTracker } from '../handlers';
|
||||
|
||||
describe('response usage aggregation', () => {
|
||||
const context: ResponseContext = {
|
||||
responseId: 'resp_test',
|
||||
model: 'agent_test',
|
||||
createdAt: 1778317637,
|
||||
};
|
||||
|
||||
it('accumulates usage across parent and subagent model calls', () => {
|
||||
const aggregator = createResponseAggregator();
|
||||
const handlers = createAggregatorEventHandlers(aggregator);
|
||||
|
|
@ -36,6 +47,56 @@ describe('response usage aggregation', () => {
|
|||
cachedTokens: 15,
|
||||
});
|
||||
});
|
||||
|
||||
it('builds one normalized wire total with an identity-free child breakdown', () => {
|
||||
const usage = buildResponsesUsage([
|
||||
{ input_tokens: 100, output_tokens: 40, provider: 'openAI' },
|
||||
{
|
||||
input_tokens: 25,
|
||||
output_tokens: 10,
|
||||
provider: 'openAI',
|
||||
usage_type: 'subagent',
|
||||
input_token_details: { cache_read: 5 },
|
||||
},
|
||||
]);
|
||||
|
||||
expect(usage).toEqual({
|
||||
input_tokens: 125,
|
||||
output_tokens: 50,
|
||||
total_tokens: 175,
|
||||
input_tokens_details: { cached_tokens: 5 },
|
||||
output_tokens_details: { reasoning_tokens: 0 },
|
||||
primary: { input_tokens: 100, output_tokens: 40, total_tokens: 140 },
|
||||
subagent: { input_tokens: 25, output_tokens: 10, total_tokens: 35 },
|
||||
});
|
||||
|
||||
const response = buildAggregatedResponse(context, createResponseAggregator(), usage);
|
||||
expect(response.usage).toEqual(usage);
|
||||
});
|
||||
|
||||
it('uses the normalized override in the completed streaming event', () => {
|
||||
const writes: string[] = [];
|
||||
const res = {
|
||||
write: (chunk: string) => {
|
||||
writes.push(chunk);
|
||||
},
|
||||
} as unknown as ServerResponse;
|
||||
const tracker = createResponseTracker();
|
||||
const usage = buildResponsesUsage([
|
||||
{ input_tokens: 100, output_tokens: 40, provider: 'openAI' },
|
||||
{
|
||||
input_tokens: 25,
|
||||
output_tokens: 10,
|
||||
provider: 'openAI',
|
||||
usage_type: 'subagent',
|
||||
},
|
||||
]);
|
||||
|
||||
createResponsesEventHandlers({ res, context, tracker }).finalizeStream(usage);
|
||||
|
||||
const completed = writes.find((chunk) => chunk.startsWith('data: {'));
|
||||
expect(JSON.parse(completed?.slice(6) ?? '{}').response.usage).toEqual(usage);
|
||||
});
|
||||
});
|
||||
|
||||
describe('convertInputToMessages', () => {
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import type {
|
|||
ReasoningTextContent,
|
||||
ItemStatus,
|
||||
ResponseStatus,
|
||||
Usage,
|
||||
} from './types';
|
||||
|
||||
/* =============================================================================
|
||||
|
|
@ -124,6 +125,7 @@ export function buildResponse(
|
|||
context: ResponseContext,
|
||||
tracker: ResponseTracker,
|
||||
status: ResponseStatus = 'in_progress',
|
||||
usageOverride?: Usage,
|
||||
): Response {
|
||||
const isCompleted = status === 'completed';
|
||||
|
||||
|
|
@ -153,13 +155,13 @@ export function buildResponse(
|
|||
reasoning: null,
|
||||
user: null,
|
||||
usage: isCompleted
|
||||
? {
|
||||
? (usageOverride ?? {
|
||||
input_tokens: tracker.usage.inputTokens,
|
||||
output_tokens: tracker.usage.outputTokens,
|
||||
total_tokens: tracker.usage.inputTokens + tracker.usage.outputTokens,
|
||||
input_tokens_details: { cached_tokens: tracker.usage.cachedTokens },
|
||||
output_tokens_details: { reasoning_tokens: tracker.usage.reasoningTokens },
|
||||
}
|
||||
})
|
||||
: null,
|
||||
max_output_tokens: null,
|
||||
max_tool_calls: null,
|
||||
|
|
@ -308,10 +310,10 @@ export function emitResponseInProgress(config: StreamHandlerConfig): void {
|
|||
/**
|
||||
* Emit response.completed event
|
||||
*/
|
||||
export function emitResponseCompleted(config: StreamHandlerConfig): void {
|
||||
export function emitResponseCompleted(config: StreamHandlerConfig, usage?: Usage): void {
|
||||
const { res, context, tracker } = config;
|
||||
tracker.status = 'completed';
|
||||
const response = buildResponse(context, tracker, 'completed');
|
||||
const response = buildResponse(context, tracker, 'completed', usage);
|
||||
writeEvent(res, {
|
||||
type: 'response.completed',
|
||||
sequence_number: tracker.nextSequence(),
|
||||
|
|
|
|||
|
|
@ -178,6 +178,7 @@ export {
|
|||
// Non-streaming
|
||||
createResponseAggregator,
|
||||
buildAggregatedResponse,
|
||||
buildResponsesUsage,
|
||||
createAggregatorEventHandlers,
|
||||
type ResponseAggregator,
|
||||
} from './service';
|
||||
|
|
|
|||
|
|
@ -13,7 +13,9 @@ import type {
|
|||
ModelContent,
|
||||
InputItem,
|
||||
Response,
|
||||
Usage,
|
||||
} from './types';
|
||||
import type { UsageMetadata } from '~/stream/interfaces/IJobStore';
|
||||
import {
|
||||
writeDone,
|
||||
emitResponseCompleted,
|
||||
|
|
@ -36,6 +38,7 @@ import {
|
|||
emitReasoningItemDone,
|
||||
type StreamHandlerConfig,
|
||||
} from './handlers';
|
||||
import { aggregateCollectedUsage } from '../usage';
|
||||
|
||||
interface ResponseUsageAccumulator {
|
||||
inputTokens: number;
|
||||
|
|
@ -361,7 +364,7 @@ interface StreamState {
|
|||
export function createResponsesEventHandlers(config: StreamHandlerConfig): {
|
||||
handlers: Record<string, { handle: (event: string, data: unknown) => void }>;
|
||||
state: StreamState;
|
||||
finalizeStream: () => void;
|
||||
finalizeStream: (usage?: Usage) => void;
|
||||
} {
|
||||
const state: StreamState = {
|
||||
messageStarted: false,
|
||||
|
|
@ -587,9 +590,9 @@ export function createResponsesEventHandlers(config: StreamHandlerConfig): {
|
|||
/**
|
||||
* Finalize the stream - close open items and emit completed
|
||||
*/
|
||||
const finalizeStream = (): void => {
|
||||
const finalizeStream = (usage?: Usage): void => {
|
||||
closeOpenStreams();
|
||||
emitResponseCompleted(config);
|
||||
emitResponseCompleted(config, usage);
|
||||
writeDone(config.res);
|
||||
};
|
||||
|
||||
|
|
@ -661,6 +664,7 @@ export function createResponseAggregator(): ResponseAggregator {
|
|||
export function buildAggregatedResponse(
|
||||
context: ResponseContext,
|
||||
aggregator: ResponseAggregator,
|
||||
usageOverride?: Usage,
|
||||
): Response {
|
||||
const output: Response['output'] = [];
|
||||
|
||||
|
|
@ -736,7 +740,7 @@ export function buildAggregatedResponse(
|
|||
top_logprobs: 0,
|
||||
reasoning: null,
|
||||
user: null,
|
||||
usage: {
|
||||
usage: usageOverride ?? {
|
||||
input_tokens: aggregator.usage.inputTokens,
|
||||
output_tokens: aggregator.usage.outputTokens,
|
||||
total_tokens: aggregator.usage.inputTokens + aggregator.usage.outputTokens,
|
||||
|
|
@ -754,6 +758,30 @@ export function buildAggregatedResponse(
|
|||
};
|
||||
}
|
||||
|
||||
/** Build provider-normalized Responses API usage from every billed call. */
|
||||
export function buildResponsesUsage(
|
||||
collectedUsage: ReadonlyArray<UsageMetadata | null | undefined>,
|
||||
): Usage {
|
||||
const { total, primary, subagent } = aggregateCollectedUsage(collectedUsage);
|
||||
return {
|
||||
input_tokens: total.inputTokens,
|
||||
output_tokens: total.outputTokens,
|
||||
total_tokens: total.totalTokens,
|
||||
input_tokens_details: { cached_tokens: total.cacheReadTokens },
|
||||
output_tokens_details: { reasoning_tokens: total.reasoningTokens },
|
||||
primary: {
|
||||
input_tokens: primary.inputTokens,
|
||||
output_tokens: primary.outputTokens,
|
||||
total_tokens: primary.totalTokens,
|
||||
},
|
||||
subagent: {
|
||||
input_tokens: subagent.inputTokens,
|
||||
output_tokens: subagent.outputTokens,
|
||||
total_tokens: subagent.totalTokens,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Create event handlers for non-streaming aggregation
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -407,6 +407,12 @@ export interface OutputTokensDetails {
|
|||
reasoning_tokens: number;
|
||||
}
|
||||
|
||||
export interface UsageTotals {
|
||||
input_tokens: number;
|
||||
output_tokens: number;
|
||||
total_tokens: number;
|
||||
}
|
||||
|
||||
/** Token usage statistics */
|
||||
export interface Usage {
|
||||
input_tokens: number;
|
||||
|
|
@ -414,6 +420,10 @@ export interface Usage {
|
|||
total_tokens: number;
|
||||
input_tokens_details: InputTokensDetails;
|
||||
output_tokens_details: OutputTokensDetails;
|
||||
/** LibreChat extension for parent, handoff, and summarization model calls. */
|
||||
primary?: UsageTotals;
|
||||
/** LibreChat extension for isolated subagent child model calls. */
|
||||
subagent?: UsageTotals;
|
||||
}
|
||||
|
||||
/** Incomplete details */
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
aggregateEmittedUsage,
|
||||
createDetachedSubagentUsageRecorder,
|
||||
createSubagentUsageSink,
|
||||
aggregateCollectedUsage,
|
||||
recordCollectedUsage,
|
||||
resolveAgentTokenConfig,
|
||||
buildPersistedContextUsage,
|
||||
|
|
@ -16,6 +17,105 @@ import {
|
|||
} from './usage';
|
||||
import { runWithDetachedSubagentUsage } from './subagentTaskContext';
|
||||
|
||||
describe('aggregateCollectedUsage', () => {
|
||||
it('preserves the no-child baseline and ignores absent entries', () => {
|
||||
expect(
|
||||
aggregateCollectedUsage([{ input_tokens: 100, output_tokens: 40, provider: 'openai' }, null]),
|
||||
).toEqual({
|
||||
total: {
|
||||
inputTokens: 100,
|
||||
outputTokens: 40,
|
||||
totalTokens: 140,
|
||||
cacheReadTokens: 0,
|
||||
reasoningTokens: 0,
|
||||
},
|
||||
primary: {
|
||||
inputTokens: 100,
|
||||
outputTokens: 40,
|
||||
totalTokens: 140,
|
||||
cacheReadTokens: 0,
|
||||
reasoningTokens: 0,
|
||||
},
|
||||
subagent: {
|
||||
inputTokens: 0,
|
||||
outputTokens: 0,
|
||||
totalTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
reasoningTokens: 0,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('includes multiple child calls once in the combined and subagent totals', () => {
|
||||
const result = aggregateCollectedUsage([
|
||||
{ input_tokens: 100, output_tokens: 40, provider: 'openai' },
|
||||
{
|
||||
input_tokens: 25,
|
||||
output_tokens: 10,
|
||||
provider: 'openai',
|
||||
usage_type: 'subagent',
|
||||
},
|
||||
{
|
||||
input_tokens: 35,
|
||||
output_tokens: 15,
|
||||
provider: 'openai',
|
||||
usage_type: 'subagent',
|
||||
},
|
||||
]);
|
||||
|
||||
expect(result.total).toEqual(
|
||||
expect.objectContaining({ inputTokens: 160, outputTokens: 65, totalTokens: 225 }),
|
||||
);
|
||||
expect(result.subagent).toEqual(
|
||||
expect.objectContaining({ inputTokens: 60, outputTokens: 25, totalTokens: 85 }),
|
||||
);
|
||||
});
|
||||
|
||||
it('uses provider-aware cache normalization for primary and child calls', () => {
|
||||
const result = aggregateCollectedUsage([
|
||||
{
|
||||
input_tokens: 200,
|
||||
output_tokens: 80,
|
||||
provider: 'anthropic',
|
||||
input_token_details: { cache_creation: 60, cache_read: 30 },
|
||||
},
|
||||
{
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
provider: 'bedrock',
|
||||
usage_type: 'subagent',
|
||||
input_token_details: { cache_creation: 20, cache_read: 10 },
|
||||
},
|
||||
]);
|
||||
|
||||
expect(result.primary.inputTokens).toBe(200);
|
||||
expect(result.subagent.inputTokens).toBe(130);
|
||||
expect(result.total.cacheReadTokens).toBe(40);
|
||||
});
|
||||
|
||||
it('repairs provider output undercounts and aggregates reasoning details', () => {
|
||||
const result = aggregateCollectedUsage([
|
||||
{
|
||||
input_tokens: 64,
|
||||
output_tokens: 2674,
|
||||
total_tokens: 3379,
|
||||
provider: 'vertexai',
|
||||
output_token_details: { reasoning: 641 },
|
||||
},
|
||||
{
|
||||
input_tokens: 20,
|
||||
output_tokens: 10,
|
||||
provider: 'openai',
|
||||
usage_type: 'subagent',
|
||||
output_token_details: { reasoning_tokens: 3 },
|
||||
},
|
||||
]);
|
||||
|
||||
expect(result.total.outputTokens).toBe(3325);
|
||||
expect(result.total.reasoningTokens).toBe(644);
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordCollectedUsage', () => {
|
||||
let mockSpendTokens: jest.Mock;
|
||||
let mockSpendStructuredTokens: jest.Mock;
|
||||
|
|
|
|||
|
|
@ -108,6 +108,20 @@ interface SplitUsage {
|
|||
completion: number;
|
||||
}
|
||||
|
||||
export interface CollectedUsageTotals {
|
||||
inputTokens: number;
|
||||
outputTokens: number;
|
||||
totalTokens: number;
|
||||
cacheReadTokens: number;
|
||||
reasoningTokens: number;
|
||||
}
|
||||
|
||||
export interface CollectedUsageBreakdown {
|
||||
total: CollectedUsageTotals;
|
||||
primary: CollectedUsageTotals;
|
||||
subagent: CollectedUsageTotals;
|
||||
}
|
||||
|
||||
function splitUsage(usage: UsageMetadata): SplitUsage {
|
||||
const cacheCreation = getCacheCreationTokens(usage);
|
||||
const cacheRead =
|
||||
|
|
@ -132,6 +146,58 @@ function splitUsage(usage: UsageMetadata): SplitUsage {
|
|||
};
|
||||
}
|
||||
|
||||
function emptyCollectedUsageTotals(): CollectedUsageTotals {
|
||||
return {
|
||||
inputTokens: 0,
|
||||
outputTokens: 0,
|
||||
totalTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
reasoningTokens: 0,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalizes every billed model call before folding it into API response totals.
|
||||
* The same provider-aware split used by billing keeps additive cache tokens and
|
||||
* repaired provider output counts consistent without coupling billing to one
|
||||
* external wire format.
|
||||
*/
|
||||
export function aggregateCollectedUsage(
|
||||
collectedUsage: ReadonlyArray<UsageMetadata | null | undefined>,
|
||||
): CollectedUsageBreakdown {
|
||||
const primary = emptyCollectedUsageTotals();
|
||||
const subagent = emptyCollectedUsageTotals();
|
||||
|
||||
for (const usage of collectedUsage) {
|
||||
if (usage == null) {
|
||||
continue;
|
||||
}
|
||||
const { totalInput, cacheRead, completion } = splitUsage(usage);
|
||||
const bucket = usage.usage_type === 'subagent' ? subagent : primary;
|
||||
const reasoningTokens =
|
||||
Number(
|
||||
usage.output_token_details?.reasoning ?? usage.output_token_details?.reasoning_tokens,
|
||||
) || 0;
|
||||
bucket.inputTokens += totalInput;
|
||||
bucket.outputTokens += completion;
|
||||
bucket.totalTokens += totalInput + completion;
|
||||
bucket.cacheReadTokens += cacheRead;
|
||||
bucket.reasoningTokens += reasoningTokens;
|
||||
}
|
||||
|
||||
return {
|
||||
total: {
|
||||
inputTokens: primary.inputTokens + subagent.inputTokens,
|
||||
outputTokens: primary.outputTokens + subagent.outputTokens,
|
||||
totalTokens: primary.totalTokens + subagent.totalTokens,
|
||||
cacheReadTokens: primary.cacheReadTokens + subagent.cacheReadTokens,
|
||||
reasoningTokens: primary.reasoningTokens + subagent.reasoningTokens,
|
||||
},
|
||||
primary,
|
||||
subagent,
|
||||
};
|
||||
}
|
||||
|
||||
export interface RecordUsageDeps {
|
||||
spendTokens: SpendTokensFn;
|
||||
spendStructuredTokens: SpendStructuredTokensFn;
|
||||
|
|
|
|||
|
|
@ -616,6 +616,8 @@ export interface UsageMetadata {
|
|||
output_token_details?: {
|
||||
/** Reasoning/thinking tokens generated as chain-of-thought (o1, Gemini thinking, etc.) */
|
||||
reasoning?: number;
|
||||
/** Alternate provider/runtime alias for reasoning tokens. */
|
||||
reasoning_tokens?: number;
|
||||
audio?: number;
|
||||
};
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue