From 4ee68d524047f69aee5625e56b3866d5a8b9a438 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 14 Jun 2026 12:00:32 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=92=B8=20feat:=20Per-Agent=20Endpoint=20T?= =?UTF-8?q?oken=20Config=20in=20Multi-Endpoint=20Billing=20(#13738)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * πŸ’Έ feat: Per-Agent Endpoint Token Config in Multi-Endpoint Billing Price each collected/emitted usage item with the producing agent's resolved endpoint token config, instead of the primary agent's for the whole graph. Previously AgentClient.recordCollectedUsage and the subagent usage emitter used a single this.options.endpointTokenConfig (the primary's) for every usage item. A connected agent or subagent on a different custom endpoint that shares a model id with an entry in the primary's tokenConfig was therefore mis-priced (a model absent from it already fell back to the built-in rate map β€” no regression). - Tag each usage with its producing agent: ModelEndHandler stamps usage.agentId = agentContext.agentId; createSubagentUsageSink stamps the child's subagentAgentId (UsageMetadata gains an optional agentId). - buildAgentToolContext retains endpointTokenConfig so initialize.js can build an agentId -> endpointTokenConfig map from agentToolContexts (the one map that holds every agent, including pure subagents pruned from agentConfigs). - AgentClient.resolveAgentEndpointTokenConfig(usage) looks up that map by agentId, falling back to the primary config; used by both the billing path (new optional resolveEndpointTokenConfig on recordCollectedUsage) and the subagent cost emitter. - recordCollectedUsage's resolver is optional and falls back to the batch endpointTokenConfig, so the shared responses.js/openai.js call sites are unchanged. - Tests: two-endpoint graph with a colliding model id prices per-agent; resolver nullish falls back to batch; subagent sink tags the child agent id. * fix: Align emit-path cost with per-agent billing; honor known-agent built-in pricing Addresses Codex review on the per-agent endpoint token config: - Emit path (callbacks.js) now prices each on_token_usage event with the producing agent's config (resolved via usageCost.resolveEndpointTokenConfig), so streamed/persisted metadata.usage.cost matches the per-agent balance transaction. The agentId tag is resolved server-side and stripped from the emitted/persisted payload. - Resolver (resolveAgentTokenConfig) now treats a known agent's config as authoritative, including undefined β†’ built-in pricing, so a known non-custom agent in a custom-primary graph is no longer charged the primary's rates. Only untagged/unknown usage falls back to the primary config. - endpointTokenConfigByAgentId records every known agent (value may be undefined) so the resolver distinguishes known-no-rates from unknown. --- .../agents/__tests__/modelEndHandler.spec.js | 39 ++++++ .../__tests__/usageEvents.integration.spec.js | 56 ++++++++ api/server/controllers/agents/callbacks.js | 18 ++- api/server/controllers/agents/client.js | 25 +++- .../services/Endpoints/agents/initialize.js | 27 ++++ .../services/Endpoints/agents/skillDeps.js | 5 + packages/api/src/agents/usage.spec.ts | 131 ++++++++++++++++++ packages/api/src/agents/usage.ts | 48 ++++++- .../api/src/stream/interfaces/IJobStore.ts | 3 + 9 files changed, 347 insertions(+), 5 deletions(-) diff --git a/api/server/controllers/agents/__tests__/modelEndHandler.spec.js b/api/server/controllers/agents/__tests__/modelEndHandler.spec.js index fdd2c88b6c..07c55e7a54 100644 --- a/api/server/controllers/agents/__tests__/modelEndHandler.spec.js +++ b/api/server/controllers/agents/__tests__/modelEndHandler.spec.js @@ -150,6 +150,45 @@ describe('ModelEndHandler β€” Vertex thoughtSignature capture (issue #13006 foll expect(collectedThoughtSignatures).toEqual({}); }); + it('tags the producing agent on collected + emitted usage for per-endpoint pricing', async () => { + const collectedUsage = []; + const emitUsage = jest.fn(); + const handler = new ModelEndHandler(collectedUsage, null, emitUsage); + const graph = { + getAgentContext: () => ({ + provider: 'openai', + agentId: 'agent_sub', + clientOptions: { model: 'gpt-4' }, + }), + }; + + await handler.handle( + 'on_chat_model_end', + { output: { usage_metadata: { input_tokens: 10, output_tokens: 5, total_tokens: 15 } } }, + { ls_model_name: 'gpt-4', run_id: 'r1', user_id: 'u1' }, + graph, + ); + + expect(collectedUsage[0].agentId).toBe('agent_sub'); + expect(emitUsage).toHaveBeenCalledWith(expect.objectContaining({ agentId: 'agent_sub' })); + }); + + it('leaves usage untagged when the graph context has no agentId (single-endpoint)', async () => { + const collectedUsage = []; + const emitUsage = jest.fn(); + const handler = new ModelEndHandler(collectedUsage, null, emitUsage); + + await handler.handle( + 'on_chat_model_end', + { output: { usage_metadata: { input_tokens: 10, output_tokens: 5, total_tokens: 15 } } }, + { ls_model_name: 'gemini-3.1-flash-lite-preview', run_id: 'r1', user_id: 'u1' }, + buildGraph(), + ); + + expect(collectedUsage[0].agentId).toBeUndefined(); + expect(emitUsage).toHaveBeenCalledWith(expect.objectContaining({ agentId: undefined })); + }); + it('throws when collectedUsage is not an array (existing contract)', () => { expect(() => new ModelEndHandler(null)).toThrow('collectedUsage must be an array'); }); diff --git a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js index 5764c159e9..9023d210cd 100644 --- a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js +++ b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js @@ -12,6 +12,7 @@ const { const { GenerationJobManager, aggregateEmittedUsage, + resolveAgentTokenConfig, buildPersistedContextUsage, } = require('@librechat/api'); const { getDefaultHandlers } = require('~/server/controllers/agents/callbacks'); @@ -301,6 +302,61 @@ describe('usage events through the real agents pipeline', () => { expect(usage.cost).toBeCloseTo(usageEmitSink.reduce((sum, e) => sum + e.cost, 0)); }); + test('emit path prices each call by its producing agent and strips the agentId tag', () => { + const res = createMockRes(); + const usageEmitSink = []; + /** Two endpoints share a model id but bill at different rates. */ + const primaryConfig = { 'gpt-4': { prompt: 0.01, completion: 0.03, context: 8192 } }; + const subagentConfig = { 'gpt-4': { prompt: 0.05, completion: 0.15, context: 8192 } }; + const byAgentId = new Map([ + ['primary', primaryConfig], + ['sub', subagentConfig], + ]); + const usageCost = { + enabled: true, + endpointTokenConfig: primaryConfig, + pricing: { + getMultiplier: ({ tokenType, model, endpointTokenConfig }) => + endpointTokenConfig?.[model]?.[tokenType] ?? 0, + getCacheMultiplier: () => 0, + }, + resolveEndpointTokenConfig: (usage) => + resolveAgentTokenConfig({ agentId: usage?.agentId, byAgentId, fallback: primaryConfig }), + }; + + const { aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res, + aggregateContent, + toolEndCallback: () => {}, + collectedUsage: [], + usageEmitSink, + usageCost, + }); + /** The CHAT_MODEL_END handler's emitUsage IS the real emitTokenUsage closure. */ + const emitUsage = handlers[GraphEvents.CHAT_MODEL_END].emitUsage; + const call = { model: 'gpt-4', input_tokens: 100, output_tokens: 50, total_tokens: 150 }; + emitUsage({ ...call, agentId: 'sub' }); + emitUsage({ ...call, agentId: 'primary' }); + + const events = res.events.filter((e) => e.event === 'on_token_usage'); + expect(events).toHaveLength(2); + /** agentId is an internal pricing tag β€” never streamed to the client nor + * folded into the persisted rollup. */ + for (const e of events) { + expect(e.data.agentId).toBeUndefined(); + } + for (const entry of usageEmitSink) { + expect(entry.agentId).toBeUndefined(); + } + /** Same tokens + model id, but the subagent endpoint's higher rates price + * its call above the primary β€” proving per-agent emit pricing. The 5Γ— ratio + * ((100Β·0.05+50Β·0.15)/(100Β·0.01+50Β·0.03)) is scale-independent of credit units. */ + expect(events[1].data.cost).toBeGreaterThan(0); + expect(events[0].data.cost).toBeGreaterThan(events[1].data.cost); + expect(events[0].data.cost / events[1].data.cost).toBeCloseTo(5); + }); + test('persists usage and context snapshot for resume via GenerationJobManager', async () => { const streamId = `usage-e2e-stream-${Date.now()}`; await GenerationJobManager.createJob(streamId, 'user-1', 'convo-1'); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index e23ffa64bd..e7cb3dfdaa 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -109,6 +109,11 @@ class ModelEndHandler { 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); /** Hidden intermediate sequential-agent calls are billed but never shown. @@ -144,6 +149,9 @@ class ModelEndHandler { model: taggedUsage.model, provider: taggedUsage.provider, usage_type: taggedUsage.usage_type, + /** Producing agent for per-endpoint pricing; consumed by the emit + * cost resolver and not included in the emitted/persisted payload. */ + agentId: taggedUsage.agentId, runId: metadata?.run_id, /** Per-run sequence so identical payloads from distinct calls * stay distinguishable during resume dedupe */ @@ -308,13 +316,19 @@ function getDefaultHandlers({ * of re-deriving from base rates. * @param {Record} data */ - const emitTokenUsage = (data) => { + const emitTokenUsage = ({ agentId, ...data }) => { let payload = data; if (usageCost?.enabled === true && usageCost.pricing) { try { + /** Price with the producing agent's config (multi-endpoint graphs) so + * the streamed/persisted cost matches the per-agent balance transaction; + * `agentId` is resolved here, not forwarded to the client or rollup. */ + const endpointTokenConfig = usageCost.resolveEndpointTokenConfig + ? usageCost.resolveEndpointTokenConfig({ agentId }) + : usageCost.endpointTokenConfig; payload = { ...data, - cost: computeUsageCostUSD(data, usageCost.pricing, usageCost.endpointTokenConfig), + cost: computeUsageCostUSD(data, usageCost.pricing, endpointTokenConfig), }; } catch (err) { logger.warn('[getDefaultHandlers] Failed to compute usage cost', err); diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 6c3d3d9c61..d73adb246f 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -24,6 +24,7 @@ const { sendEvent, computeUsageCostUSD, aggregateEmittedUsage, + resolveAgentTokenConfig, buildPersistedContextUsage, createSubagentUsageSink, isDeepSeekReasoningProvider, @@ -886,6 +887,24 @@ class AgentClient extends BaseClient { return Object.keys(metadata).length > 0 ? metadata : undefined; } + /** + * Resolves the endpoint token config for a usage item by its producing agent + * (multi-endpoint graphs: connected agents + subagents). A known agent's + * config is authoritative β€” including `undefined`, which prices with built-in + * rates (e.g. a non-custom agent in a custom-primary graph). Only an + * untagged/unknown agent falls back to the primary config, so single-endpoint + * graphs are unchanged. + * @param {UsageMetadata} usage + * @returns {import('@librechat/api').EndpointTokenConfig | undefined} + */ + resolveAgentEndpointTokenConfig(usage) { + return resolveAgentTokenConfig({ + agentId: usage?.agentId, + byAgentId: this.options.endpointTokenConfigByAgentId, + fallback: this.options.endpointTokenConfig, + }); + } + /** * @param {Object} params * @param {string} [params.model] @@ -918,6 +937,7 @@ class AgentClient extends BaseClient { balance, transactions, endpointTokenConfig: this.options.endpointTokenConfig, + resolveEndpointTokenConfig: (usage) => this.resolveAgentEndpointTokenConfig(usage), }, ); @@ -950,7 +970,6 @@ class AgentClient extends BaseClient { return undefined; } const includeCost = appConfig?.interfaceConfig?.contextCost === true; - const endpointTokenConfig = this.options.endpointTokenConfig; return (usage) => { const data = { input_tokens: usage.input_tokens, @@ -963,11 +982,13 @@ class AgentClient extends BaseClient { runId: this.responseMessageId, /** Unique per collected entry (post-push length) for resume dedupe */ seq: this.collectedUsage.length, + /** Price with the SUBAGENT's own endpoint token config (its endpoint may + * differ from the parent's); `usage.agentId` is tagged by the sink. */ cost: includeCost ? computeUsageCostUSD( usage, { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier }, - endpointTokenConfig, + this.resolveAgentEndpointTokenConfig(usage), ) : undefined, }; diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index f72ef653e9..9caee651c4 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -9,6 +9,7 @@ const { GenerationJobManager, getCustomEndpointConfig, discoverConnectedAgents, + resolveAgentTokenConfig, resolveAgentScopedSkillIds, resolveModelSpecSkillIds, buildAgentContextAttachmentsByAgentId, @@ -904,6 +905,28 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { }) : undefined; + /** Per-agent resolved endpoint token config, keyed by agent id. Built from + * `agentToolContexts` (the one map holding every agent, including pure + * subagents pruned from `agentConfigs`) so usage billed/emitted for a + * connected or subagent on a different custom endpoint is priced with THAT + * agent's configured rates instead of the primary's. Every known agent is + * recorded β€” even with an `undefined` config β€” so the resolver can tell a + * known non-custom agent (built-in pricing) from an untagged/unknown one + * (primary fallback). + * @type {Map} */ + const endpointTokenConfigByAgentId = new Map(); + for (const [agentId, ctx] of agentToolContexts) { + endpointTokenConfigByAgentId.set(agentId, ctx?.endpointTokenConfig); + } + /** Price emitted usage per producing agent too, so the streamed/persisted + * `metadata.usage.cost` matches the per-agent balance transaction. */ + usageCost.resolveEndpointTokenConfig = (usage) => + resolveAgentTokenConfig({ + agentId: usage?.agentId, + byAgentId: endpointTokenConfigByAgentId, + fallback: usageCost.endpointTokenConfig, + }); + const client = new AgentClient({ req, res, @@ -930,6 +953,10 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { /** Resolved endpoint token/pricing config so spending and cost reflect * configured rates for custom-endpoint agents instead of defaults. */ endpointTokenConfig: primaryConfig.endpointTokenConfig, + /** Per-agent override of the above for multi-endpoint graphs (connected + * agents + subagents); falls back to the primary config when an agent + * isn't present or has no configured rates. */ + endpointTokenConfigByAgentId, /** Capture sinks the handlers fill during the run; `sendCompletion` reads * them to persist the breakdown + usage rollup on the response message. */ contextUsageSink, diff --git a/api/server/services/Endpoints/agents/skillDeps.js b/api/server/services/Endpoints/agents/skillDeps.js index 83d287fc08..7154c1d52b 100644 --- a/api/server/services/Endpoints/agents/skillDeps.js +++ b/api/server/services/Endpoints/agents/skillDeps.js @@ -272,6 +272,11 @@ function buildSkillPrimedIdsByName(manualSkillPrimes, alwaysApplySkillPrimes) { function buildAgentToolContext({ agent, config }) { return { agent, + /** Per-agent resolved endpoint token/pricing config. Retained here because + * `agentToolContexts` is the one map that holds every agent β€” including + * pure subagents pruned from `agentConfigs` β€” so usage can be priced with + * the producing agent's config in multi-endpoint graphs. */ + endpointTokenConfig: config.endpointTokenConfig, toolRegistry: config.toolRegistry, mcpAvailableTools: config.mcpAvailableTools, requestScopedConnections: config.requestScopedConnections, diff --git a/packages/api/src/agents/usage.spec.ts b/packages/api/src/agents/usage.spec.ts index 8f4bd96b7e..732db1da19 100644 --- a/packages/api/src/agents/usage.spec.ts +++ b/packages/api/src/agents/usage.spec.ts @@ -7,6 +7,7 @@ import { aggregateEmittedUsage, createSubagentUsageSink, recordCollectedUsage, + resolveAgentTokenConfig, buildPersistedContextUsage, buildAbortedResponseMetadata, } from './usage'; @@ -809,6 +810,77 @@ describe('recordCollectedUsage', () => { ); }); + it('prices each usage with its agent endpoint config (multi-endpoint graph)', async () => { + /** Two endpoints share a model id but bill at different rates. */ + const primaryConfig = { 'shared-model': { prompt: 0.01, completion: 0.03, context: 8192 } }; + const subagentConfig = { 'shared-model': { prompt: 0.05, completion: 0.15, context: 8192 } }; + const byAgent: Record = { + primary: primaryConfig, + sub: subagentConfig, + }; + const collectedUsage: UsageMetadata[] = [ + { + usage_type: 'message', + agentId: 'primary', + input_tokens: 100, + output_tokens: 50, + model: 'shared-model', + }, + { + usage_type: 'subagent', + agentId: 'sub', + input_tokens: 200, + output_tokens: 80, + model: 'shared-model', + }, + ]; + + await recordCollectedUsage(deps, { + ...baseParams, + collectedUsage, + endpointTokenConfig: primaryConfig, + resolveEndpointTokenConfig: (usage) => + usage.agentId != null ? byAgent[usage.agentId] : undefined, + }); + + expect(mockSpendTokens).toHaveBeenCalledWith( + expect.objectContaining({ + context: 'message', + model: 'shared-model', + endpointTokenConfig: primaryConfig, + }), + { promptTokens: 100, completionTokens: 50 }, + ); + expect(mockSpendTokens).toHaveBeenCalledWith( + expect.objectContaining({ + context: 'subagent', + model: 'shared-model', + endpointTokenConfig: subagentConfig, + }), + { promptTokens: 200, completionTokens: 80 }, + ); + }); + + it('trusts the resolver result, including undefined (built-in pricing for known agents)', async () => { + const batch = { 'gpt-4': { prompt: 0.01, completion: 0.03, context: 8192 } }; + await recordCollectedUsage(deps, { + ...baseParams, + collectedUsage: [ + { agentId: 'known-openai', input_tokens: 100, output_tokens: 50, model: 'gpt-4' }, + ], + endpointTokenConfig: batch, + /** A known agent with no configured rates: the resolver returns undefined + * and the billing path must honor it (built-in pricing), NOT re-apply + * the batch/primary config. */ + resolveEndpointTokenConfig: () => undefined, + }); + + expect(mockSpendTokens).toHaveBeenCalledWith( + expect.objectContaining({ endpointTokenConfig: undefined }), + { promptTokens: 100, completionTokens: 50 }, + ); + }); + it('should use default context "message" when not provided', async () => { const collectedUsage: UsageMetadata[] = [ { input_tokens: 100, output_tokens: 50, model: 'gpt-4' }, @@ -1382,10 +1454,24 @@ describe('createSubagentUsageSink', () => { total_tokens: 1600, model: 'claude-haiku-4-5', provider: 'anthropic', + agentId: 'researcher', }, ]); }); + it('tags the child agent id so the host can price with the subagent endpoint config', () => { + const collectedUsage: UsageMetadata[] = []; + const emitted: UsageMetadata[] = []; + const sink = createSubagentUsageSink(collectedUsage, (u) => emitted.push(u)); + + sink(makeEvent({ subagentAgentId: 'agent_xyz' })); + + expect(collectedUsage[0].agentId).toBe('agent_xyz'); + /** The same tagged object is handed to onUsage (the live emitter). */ + expect(emitted[0]).toBe(collectedUsage[0]); + expect(emitted[0].agentId).toBe('agent_xyz'); + }); + it('preserves cache token details from the child call', () => { const collectedUsage: UsageMetadata[] = []; const sink = createSubagentUsageSink(collectedUsage); @@ -1698,3 +1784,48 @@ describe('buildAbortedResponseMetadata', () => { expect((result as { contextUsage?: unknown }).contextUsage).toBeUndefined(); }); }); + +describe('resolveAgentTokenConfig', () => { + const primary = { 'gpt-4': { prompt: 0.01, completion: 0.03, context: 8192 } }; + const subagent = { 'gpt-4': { prompt: 0.05, completion: 0.15, context: 8192 } }; + + it('returns the producing agent’s own config', () => { + const byAgentId = new Map([ + ['primary', primary], + ['sub', subagent], + ]); + expect(resolveAgentTokenConfig({ agentId: 'sub', byAgentId, fallback: primary })).toBe( + subagent, + ); + }); + + it('returns undefined for a known agent with no configured rates (built-in pricing)', () => { + /** A known non-custom agent (e.g. a normal OpenAI agent) is recorded with an + * undefined config; it must NOT inherit the custom-primary rates. */ + const byAgentId = new Map([ + ['primary', primary], + ['known-openai', undefined], + ]); + expect( + resolveAgentTokenConfig({ agentId: 'known-openai', byAgentId, fallback: primary }), + ).toBeUndefined(); + }); + + it('falls back to the primary config for an untagged usage', () => { + const byAgentId = new Map([['primary', primary]]); + expect(resolveAgentTokenConfig({ agentId: undefined, byAgentId, fallback: primary })).toBe( + primary, + ); + }); + + it('falls back to the primary config for an unknown agent id', () => { + const byAgentId = new Map([['primary', primary]]); + expect(resolveAgentTokenConfig({ agentId: 'ghost', byAgentId, fallback: primary })).toBe( + primary, + ); + }); + + it('returns the fallback when there is no per-agent map (single-endpoint graphs)', () => { + expect(resolveAgentTokenConfig({ agentId: 'primary', fallback: primary })).toBe(primary); + }); +}); diff --git a/packages/api/src/agents/usage.ts b/packages/api/src/agents/usage.ts index 43ff590a86..e8b4583cff 100644 --- a/packages/api/src/agents/usage.ts +++ b/packages/api/src/agents/usage.ts @@ -326,6 +326,31 @@ export function buildAbortedResponseMetadata( return usage ? { usage } : undefined; } +/** + * Resolves the endpoint token config for a usage item by its producing agent. + * Multi-endpoint graphs tag each call with `agentId`; that agent's resolved + * config is authoritative β€” including `undefined`, which means "no configured + * rates, use built-in pricing" (e.g. a non-custom agent in a custom-primary + * graph). Only an untagged or unknown agent falls back to `fallback` (the + * primary config), so single-endpoint graphs are unchanged. `byAgentId` must + * hold an entry for every known agent (value may be `undefined`) so `has` + * distinguishes "known, no rates" from "unknown". + */ +export function resolveAgentTokenConfig({ + agentId, + byAgentId, + fallback, +}: { + agentId?: string | null; + byAgentId?: Map; + fallback?: EndpointTokenConfig; +}): EndpointTokenConfig | undefined { + if (agentId != null && byAgentId?.has(agentId)) { + return byAgentId.get(agentId); + } + return fallback; +} + export interface RecordUsageParams { user: string; conversationId: string; @@ -336,6 +361,15 @@ export interface RecordUsageParams { balance?: Partial | null; transactions?: Partial; endpointTokenConfig?: EndpointTokenConfig; + /** + * Per-usage endpoint token config resolver for multi-endpoint graphs. Called + * with each usage item; when provided it is authoritative β€” its result prices + * that item, including `undefined` (built-in pricing for a known agent with no + * configured rates). It owns its own fallback to the primary config for + * untagged/unknown agents (see {@link resolveAgentTokenConfig}). Single-config + * callers (responses.js / openai.js) omit it and use `endpointTokenConfig`. + */ + resolveEndpointTokenConfig?: (usage: UsageMetadata) => EndpointTokenConfig | undefined; } export interface RecordUsageResult { @@ -363,6 +397,7 @@ export async function recordCollectedUsage( conversationId, collectedUsage, endpointTokenConfig, + resolveEndpointTokenConfig, context = 'message', } = params; @@ -422,7 +457,12 @@ export async function recordCollectedUsage( messageId, transactions, conversationId, - endpointTokenConfig, + /** Price with the producing agent's endpoint config when a resolver is + * provided (multi-endpoint graphs); it owns the fallback to the primary + * config, so `undefined` here means built-in pricing, not the batch one. */ + endpointTokenConfig: resolveEndpointTokenConfig + ? resolveEndpointTokenConfig(usage) + : endpointTokenConfig, context: usageContext, model: usage.model ?? model, }; @@ -562,6 +602,12 @@ export function createSubagentUsageSink( if (event.provider != null && event.provider !== '') { usage.provider = event.provider; } + /** Tag the child's agent id so the host can price this usage with the + * subagent's own endpoint token config (its endpoint may differ from the + * parent's). The same tagged object is pushed AND handed to `onUsage`. */ + if (event.subagentAgentId != null && event.subagentAgentId !== '') { + usage.agentId = event.subagentAgentId; + } collectedUsage.push(usage); /** Lets the host stream the billed child usage to the client (tagged * `subagent`, so it folds into session cost/totals but not the live diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 32b42338f0..fcb472c7cd 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -93,6 +93,9 @@ export interface UsageMetadata { model?: string; /** Provider identifier that generated this usage */ provider?: string; + /** Agent that produced this usage (graph agent id / subagent agent id). Lets + * multi-endpoint graphs price each call with its own endpoint token config. */ + agentId?: string; /** * OpenAI-style cache token details. * Present for OpenAI models (GPT-4, o1, etc.)