diff --git a/backend/__tests__/unit/routes/runTraces.test.js b/backend/__tests__/unit/routes/runTraces.test.js index 3769f4fc83..94983bcc09 100644 --- a/backend/__tests__/unit/routes/runTraces.test.js +++ b/backend/__tests__/unit/routes/runTraces.test.js @@ -65,6 +65,20 @@ test('returns durable turn + PG attribution after card retention without content expect(JSON.stringify(res.body)).not.toMatch(/PRIVATE_|private prompt|argsDigest/); expect(res.body.retention).toEqual({ cardDefaultDays: 30, traceExpiresAt: null }); }); +test('a turn keeps the model that answered beside the one requested, through save and read (TASK-028)', async () => { + run.turns = [ + { turnIndex: 0, model: 'openrouter/nvidia/nemotron-3-super-120b-a12b:free', servedModel: 'deepseek-v4-flash', fallbacks: 1, elapsedMs: 1 }, + { turnIndex: 1, model: 'deepseek-v4-flash', servedModel: 'deepseek-v4-flash', elapsedMs: 1 }, + { turnIndex: 2, model: 'model', elapsedMs: 1 }, + ]; + await run.save(); + const res = await get(); + expect(res.body.run.turns.map(({ model, servedModel, fallbacks }) => ({ model, servedModel, fallbacks }))).toEqual([ + { model: 'openrouter/nvidia/nemotron-3-super-120b-a12b:free', servedModel: 'deepseek-v4-flash', fallbacks: 1 }, + { model: 'deepseek-v4-flash', servedModel: 'deepseek-v4-flash' }, + { model: 'model' }, + ]); +}); test.each(['global admin', 'agent-dm observer'])( 'returns trace metadata without authorization identities to a non-member %s', async (viewerType) => { if (viewerType === 'global admin') { diff --git a/backend/__tests__/unit/services/nativeRuntimeService.servedModel.test.js b/backend/__tests__/unit/services/nativeRuntimeService.servedModel.test.js new file mode 100644 index 0000000000..268d265a47 --- /dev/null +++ b/backend/__tests__/unit/services/nativeRuntimeService.servedModel.test.js @@ -0,0 +1,136 @@ +/** + * TASK-028: a native turn records the model group that answered it beside the + * one it requested. LiteLLM names the answering group in x-litellm-model-group + * and counts fallback hops in x-litellm-attempted-fallbacks (v1.106.2); the + * body's `model` is the second source. + * + * Real HTTP: axios is not mocked. A local server stands in for the LiteLLM + * proxy and answers with the headers the proxy sends, so the header names and + * their parsing are exercised as they arrive over the wire. + */ +const http = require('http'); + +jest.mock('../../../models/Pod', () => ({ + findById: jest.fn(() => ({ + lean: jest.fn().mockResolvedValue({ name: 'Served pod' }), + })), +})); + +jest.mock('../../../models/AgentRun', () => ({ + create: jest.fn(), + countDocuments: jest.fn(), +})); + +jest.mock('../../../services/agentTypingService', () => ({ + emitAgentTypingStart: jest.fn(), + emitAgentTypingStop: jest.fn(), +})); + +jest.mock('../../../services/agentMessageService', () => ({ + postMessage: jest.fn().mockResolvedValue({ success: true }), +})); + +jest.mock('../../../services/messageClaimService', () => ({ + claim: jest.fn(), + release: jest.fn(), +})); + +const AgentRun = require('../../../models/AgentRun'); +const MessageClaimService = require('../../../services/messageClaimService'); +const { runAgent } = require('../../../services/nativeRuntimeService'); + +const NEMOTRON = 'openrouter/nvidia/nemotron-3-super-120b-a12b:free'; +const DEEPSEEK = 'deepseek-v4-flash'; + +let server; +let reply; +let requests; +let savedRun; +const previousEnv = { ...process.env }; + +beforeAll(async () => { + server = http.createServer((req, res) => { + let body = ''; + req.on('data', (chunk) => { body += chunk; }); + req.on('end', () => { + requests.push({ url: req.url, body: JSON.parse(body || '{}') }); + res.writeHead(200, { 'content-type': 'application/json', ...reply.headers }); + res.end(JSON.stringify({ + ...(reply.model ? { model: reply.model } : {}), + choices: [{ message: { role: 'assistant', content: 'NO_REPLY' } }], + usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 }, + })); + }); + }); + await new Promise((resolve) => { server.listen(0, '127.0.0.1', resolve); }); + process.env.LITELLM_BASE_URL = `http://127.0.0.1:${server.address().port}`; + process.env.LITELLM_API_KEY = 'test-service-key'; +}); + +afterAll(async () => { + process.env = previousEnv; + await new Promise((resolve) => { server.close(resolve); }); +}); + +beforeEach(() => { + jest.clearAllMocks(); + MessageClaimService.claim.mockResolvedValue({ claimed: true, expiresAt: new Date() }); + MessageClaimService.release.mockResolvedValue({ released: true }); + AgentRun.create.mockImplementation(async (doc) => { + savedRun = { _id: 'run-1', ...doc, save: jest.fn().mockResolvedValue(undefined) }; + return savedRun; + }); +}); + +const installation = (model) => ({ + podId: 'pod-1', + agentName: 'served', + instanceId: 'default', + displayName: 'Served', + config: { + runtime: { runtimeType: 'native' }, systemPrompt: 'answer', model, tools: [], + }, +}); + +const mention = { + type: 'chat.mention', + eventId: 'evt-1', + payload: { messageId: 'msg-1', content: 'hi @served' }, +}; + +const turnOf = async (model, nextReply) => { + reply = nextReply; + requests = []; + const result = await runAgent(installation(model), mention); + expect(result.status).toBe('succeeded'); + expect(requests).toHaveLength(1); + const [turn] = savedRun.turns; + return { model: turn.model, servedModel: turn.servedModel, fallbacks: turn.fallbacks }; +}; + +test('a turn another group answered records both, and the hops', async () => { + const turn = await turnOf(NEMOTRON, { + // Through the proxy a fallback keeps the provider's own name in the body. + model: 'deepseek/deepseek-chat', + headers: { 'x-litellm-model-group': DEEPSEEK, 'x-litellm-attempted-fallbacks': '1' }, + }); + expect(turn).toEqual({ model: NEMOTRON, servedModel: DEEPSEEK, fallbacks: 1 }); + // The refusal still goes out: recording a fallback is not permission for one (D81). + expect(requests[0].body).toMatchObject({ model: NEMOTRON, disable_fallbacks: true }); +}); + +test('a normal turn records the two equal, with no hops', async () => { + const turn = await turnOf(DEEPSEEK, { + model: DEEPSEEK, + headers: { 'x-litellm-model-group': DEEPSEEK, 'x-litellm-attempted-fallbacks': '0' }, + }); + expect(turn).toEqual({ model: DEEPSEEK, servedModel: DEEPSEEK, fallbacks: undefined }); +}); + +test('without the group header the body model is used, and without either nothing is recorded', async () => { + expect(await turnOf(NEMOTRON, { model: DEEPSEEK, headers: {} })) + .toEqual({ model: NEMOTRON, servedModel: DEEPSEEK, fallbacks: undefined }); + // An unmeasured turn is not recorded as served by the requested model. + expect(await turnOf(NEMOTRON, { headers: { 'x-litellm-attempted-fallbacks': 'not-a-number' } })) + .toEqual({ model: NEMOTRON, servedModel: undefined, fallbacks: undefined }); +}); diff --git a/backend/__tests__/unit/services/spendAlarmService.test.js b/backend/__tests__/unit/services/spendAlarmService.test.js index 9604f6d265..02be147b79 100644 --- a/backend/__tests__/unit/services/spendAlarmService.test.js +++ b/backend/__tests__/unit/services/spendAlarmService.test.js @@ -84,6 +84,20 @@ describe('userSpendForDay', () => { expect(rows.find((r) => r.userId === String(USER_A)).usd).toBeCloseTo(spend.UNPRICED_MODEL.input, 6); expect(rows.find((r) => r.userId === String(USER_B)).usd).toBeCloseTo(1.2, 6); }); + + test('a turn is priced by the model that answered it, and an older turn by the one it asked for', async () => { + await AgentRun.create([ + // Asked DeepSeek, answered by Sonnet: 1M out at Sonnet's $10, not DeepSeek's $1.20 (TASK-028). + run(USER_A, [{ model: 'deepseek-v4-flash', servedModel: 'claude-sonnet-5-5', fallbacks: 1, + promptTokens: 0, completionTokens: 1_000_000 }]), + // Recorded before servedModel existed: priced by model, as before. + run(USER_B, [{ model: 'claude-sonnet-5-5', promptTokens: 0, completionTokens: 1_000_000 }]), + ]); + + const rows = await spend.userSpendForDay(NOW); + expect(rows.find((r) => r.userId === String(USER_A)).usd).toBeCloseTo(10, 6); + expect(rows.find((r) => r.userId === String(USER_B)).usd).toBeCloseTo(10, 6); + }); }); describe('sandbox seconds (TASK-266)', () => { diff --git a/backend/models/AgentRun.ts b/backend/models/AgentRun.ts index 96ad06ea03..7ef6ea60dd 100644 --- a/backend/models/AgentRun.ts +++ b/backend/models/AgentRun.ts @@ -65,6 +65,13 @@ export interface IAgentRunTurn { llmResponseText?: string; elapsedMs: number; liteLLMCallId?: string; + // The model group that answered, when it can differ from `model`, the group + // requested (TASK-028). Absent on turns recorded before it was kept, and on a + // turn whose response named no model. + servedModel?: string; + // LiteLLM's x-litellm-attempted-fallbacks; absent means none. The proxy reports + // the outermost hop's depth, so a two-hop chain records 1: a fallback happened. + fallbacks?: number; } export interface IAgentRun extends Document { @@ -132,6 +139,8 @@ const AgentRunTurnSchema = new Schema( llmResponseText: { type: String }, elapsedMs: { type: Number, default: 0 }, liteLLMCallId: { type: String }, + servedModel: { type: String }, + fallbacks: { type: Number }, }, { _id: false }, ); diff --git a/backend/routes/runTraces.ts b/backend/routes/runTraces.ts index b45e29efea..b5f9b8f374 100644 --- a/backend/routes/runTraces.ts +++ b/backend/routes/runTraces.ts @@ -37,7 +37,8 @@ router.get('/:podId/runs/:runId', traceReadLimit, auth, async (req, res) => { .select([ 'agentName instanceId status trigger startedAt completedAt totalTokens errorKind', 'traceTruncated omittedTurns omittedToolCalls', - 'turns.turnIndex turns.model turns.elapsedMs turns.promptTokens turns.completionTokens turns.totalTokens', + 'turns.turnIndex turns.model turns.servedModel turns.fallbacks turns.elapsedMs', + 'turns.promptTokens turns.completionTokens turns.totalTokens', 'turns.toolCalls.name turns.toolCalls.elapsedMs turns.toolCalls.callId turns.toolCalls.outcome', ].join(' ')) .lean(); @@ -59,7 +60,8 @@ router.get('/:podId/runs/:runId', traceReadLimit, auth, async (req, res) => { traceTruncated: Boolean(run.traceTruncated), omittedTurns: run.omittedTurns || 0, omittedToolCalls: run.omittedToolCalls || 0, turns: run.turns.map((turn) => ({ - turnIndex: turn.turnIndex, model: turn.model, elapsedMs: turn.elapsedMs, + turnIndex: turn.turnIndex, model: turn.model, servedModel: turn.servedModel, + fallbacks: turn.fallbacks, elapsedMs: turn.elapsedMs, promptTokens: turn.promptTokens, completionTokens: turn.completionTokens, totalTokens: turn.totalTokens, toolCalls: turn.toolCalls.map((call) => ({ name: call.name, elapsedMs: call.elapsedMs, callId: call.callId, diff --git a/backend/services/agentRunLimits.ts b/backend/services/agentRunLimits.ts index 29b061fd15..c1fda2da0a 100644 --- a/backend/services/agentRunLimits.ts +++ b/backend/services/agentRunLimits.ts @@ -45,7 +45,8 @@ export function boundAgentRun(run: { turnIndex: original.turnIndex, model: text(original.model) || 'unknown', promptTokens: original.promptTokens, completionTokens: original.completionTokens, totalTokens: original.totalTokens, elapsedMs: original.elapsedMs, - liteLLMCallId: text(original.liteLLMCallId), + liteLLMCallId: text(original.liteLLMCallId), servedModel: text(original.servedModel), + fallbacks: original.fallbacks, llmResponseText: text(original.llmResponseText, 4096), toolCalls: [], }; if (original.llmResponseText !== turn.llmResponseText) truncated = true; diff --git a/backend/services/nativeRuntimeService.ts b/backend/services/nativeRuntimeService.ts index 4ec7e856d8..c95643fc76 100644 --- a/backend/services/nativeRuntimeService.ts +++ b/backend/services/nativeRuntimeService.ts @@ -1395,6 +1395,7 @@ interface LiteLLMChoice { } interface LiteLLMResponse { + model?: string; choices?: LiteLLMChoice[]; usage?: { prompt_tokens?: number; @@ -1413,6 +1414,29 @@ function extractCallId(data: LiteLLMResponse, headers: Record): return headerId ? String(headerId) : undefined; } +// Which model group answered the call (TASK-028). LiteLLM's router names it in +// x-litellm-model-group on every response, and x-litellm-attempted-fallbacks +// says a fallback happened (v1.106.2; it is the outermost hop's depth, so 1 on a +// two-hop chain). The proxy restamps the body's `model` to the requested group +// unless a fallback happened, so the body comes second. When neither names a +// model, nothing was measured and the field stays unset; every reader falls +// back to the requested model. Native calls send disable_fallbacks (D81), so +// the two differ only if something overrode it, such as a key whose metadata +// sets disable_fallbacks. +function servedModelOf( + data: LiteLLMResponse, + headers: Record, +): { servedModel?: string; fallbacks?: number } { + const group = headers?.['x-litellm-model-group']; + const hops = Number(headers?.['x-litellm-attempted-fallbacks']); + const servedModel = (typeof group === 'string' && group) + || (typeof data?.model === 'string' && data.model) || undefined; + return { + ...(servedModel ? { servedModel } : {}), + ...(Number.isInteger(hops) && hops > 0 ? { fallbacks: hops } : {}), + }; +} + function mapTriggerType(raw: string): NativeRunTrigger['type'] { if (raw === 'chat.mention' || raw === 'thread.mention' || raw === 'mention') return 'mention'; if (raw === 'heartbeat') return 'heartbeat'; @@ -2142,6 +2166,7 @@ export async function runAgent( llmResponseText: (msg.content as string | null) || undefined, elapsedMs: 0, liteLLMCallId: extractCallId(llmResponse, responseHeaders), + ...servedModelOf(llmResponse, responseHeaders), }; const toolCalls = msg.tool_calls || []; diff --git a/backend/services/spendAlarmService.ts b/backend/services/spendAlarmService.ts index a7f58bd336..1bad45ff0c 100644 --- a/backend/services/spendAlarmService.ts +++ b/backend/services/spendAlarmService.ts @@ -95,7 +95,9 @@ export const userSpendForDay = async (now: Date = new Date()): Promise { + await act(async () => { await i18n.changeLanguage(language); }); + const nemotron = 'openrouter/nvidia/nemotron-3-super-120b-a12b:free'; + mockGet.mockResolvedValue({ run: { ...trace, turns: [ + { ...trace.turns[0], model: nemotron, servedModel: 'deepseek-v4-flash', fallbacks: 1 }, + { ...trace.turns[0], turnIndex: 1, model: 'deepseek-v4-flash', servedModel: 'deepseek-v4-flash' }, + { ...trace.turns[0], turnIndex: 2, model: 'older-turn' }, + ] } }); + const { container } = setup(); + await screen.findAllByText('vendor.read'); + const meta = Array.from(container.querySelectorAll('.v2-run-trace__meta'), (node) => node.textContent); + expect(meta[0]).toBe(`${i18n.t('runTrace.served', { requested: nemotron, served: 'deepseek-v4-flash' })} · 12 ms · 123 ${i18n.t('runTrace.tokens')}`); + expect(meta[0]).toContain(nemotron); + expect(meta[0]).toContain('deepseek-v4-flash'); + expect(meta[1]).toBe(`deepseek-v4-flash · 12 ms · 123 ${i18n.t('runTrace.tokens')}`); + expect(meta[2]).toBe(`older-turn · 12 ms · 123 ${i18n.t('runTrace.tokens')}`); +}); test('refusal can retry and never presents an empty successful trace', async () => { mockGet.mockRejectedValueOnce({ response: { status: 404 } }).mockResolvedValueOnce({ run: trace }); setup(); diff --git a/frontend/src/v2/components/V2RunTrace.tsx b/frontend/src/v2/components/V2RunTrace.tsx index 72182c1a3d..b72eb22bba 100644 --- a/frontend/src/v2/components/V2RunTrace.tsx +++ b/frontend/src/v2/components/V2RunTrace.tsx @@ -9,7 +9,8 @@ export interface RunTrace { trigger: string; startedAt: string; completedAt?: string; totalTokens: number; errorKind?: string; traceTruncated: boolean; omittedTurns: number; omittedToolCalls: number; turns: Array<{ - turnIndex: number; model: string; elapsedMs: number; totalTokens?: number; + turnIndex: number; model: string; servedModel?: string; fallbacks?: number; + elapsedMs: number; totalTokens?: number; toolCalls: Array<{ name: string; elapsedMs: number; callId?: string; outcome: string; trail?: { grantId?: string; credentialOwnerId?: string | null; outcome: string } | null; @@ -61,7 +62,9 @@ const V2RunTrace: React.FC = () => {
    {trace.turns.map((turn) =>
  1. {t('runTrace.turn', { number: turn.turnIndex + 1 })}

    -

    {turn.model} · {turn.elapsedMs} ms · {turn.totalTokens ?? '—'} {t('runTrace.tokens')}

    +

    {turn.servedModel && turn.servedModel !== turn.model + ? t('runTrace.served', { requested: turn.model, served: turn.servedModel }) + : turn.model} · {turn.elapsedMs} ms · {turn.totalTokens ?? '—'} {t('runTrace.tokens')}

    {!turn.toolCalls.length &&

    {t('runTrace.noTools')}

    } {turn.toolCalls.map((call, index) =>
    {call.name}{enumLabel('outcomes', call.outcome)} · {call.elapsedMs} ms