diff --git a/agentic/agentic-server/__tests__/gateway.test.ts b/agentic/agentic-server/__tests__/gateway.test.ts index 1292cc9a89..276b855b75 100644 --- a/agentic/agentic-server/__tests__/gateway.test.ts +++ b/agentic/agentic-server/__tests__/gateway.test.ts @@ -142,6 +142,56 @@ describe('POST /v1/chat/completions', () => { }); }); + it('forwards task-correlation headers into the sink entry so tokens join to the task', async () => { + await fetch(`http://localhost:${agenticPort}/v1/chat/completions`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Database-Id': 'db-link-1', + 'X-Entity-Id': 'entity-link-1', + 'X-Actor-Id': 'actor-link-1', + 'X-Invocation-Id': '0190a5f0-0000-7000-8000-000000000001', + 'X-Job-Id': '4242', + 'X-Attempt': '2', + 'X-Run-Id': '0190a5f0-0000-7000-8000-000000000002' + }, + body: JSON.stringify({ model: 'llama3', messages: [{ role: 'user', content: 'linked' }] }) + }); + + await new Promise((r) => setTimeout(r, 100)); + + expect(sinkEntries).toHaveLength(1); + expect(sinkEntries[0]).toMatchObject({ + databaseId: 'db-link-1', + entityId: 'entity-link-1', + actorId: 'actor-link-1', + invocationId: '0190a5f0-0000-7000-8000-000000000001', + jobId: '4242', + attempt: 2, + runId: '0190a5f0-0000-7000-8000-000000000002' + }); + }); + + it('leaves linkage undefined when no correlation headers arrive and drops a malformed X-Attempt', async () => { + await fetch(`http://localhost:${agenticPort}/v1/chat/completions`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Database-Id': 'db-unlinked-1', + 'X-Attempt': 'two' + }, + body: JSON.stringify({ model: 'llama3', messages: [{ role: 'user', content: 'unlinked' }] }) + }); + + await new Promise((r) => setTimeout(r, 100)); + + expect(sinkEntries).toHaveLength(1); + expect(sinkEntries[0].invocationId).toBeUndefined(); + expect(sinkEntries[0].jobId).toBeUndefined(); + expect(sinkEntries[0].attempt).toBeUndefined(); + expect(sinkEntries[0].runId).toBeUndefined(); + }); + it('flattens content parts into the single string ollama accepts', async () => { // pi and every other harness send `content` as parts; ollama's chat api takes // only a string, so the gateway is where the dialects meet. @@ -273,6 +323,26 @@ describe('header security (isPublic)', () => { await new Promise((r) => setTimeout(r, 50)); expect(publicSinkEntries).toHaveLength(0); }); + + it('strips task-correlation headers too, so an external client cannot pin usage onto a task', async () => { + publicSinkEntries.length = 0; + const res = await fetch(`http://localhost:${publicPort}/v1/usage`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Database-Id': 'SHOULD-BE-STRIPPED', + 'X-Invocation-Id': 'SHOULD-BE-STRIPPED', + 'X-Job-Id': '1', + 'X-Attempt': '1', + 'X-Run-Id': 'SHOULD-BE-STRIPPED' + }, + body: JSON.stringify({ model: 'm', total_tokens: 1 }) + }); + + expect(res.status).toBe(400); + await new Promise((r) => setTimeout(r, 50)); + expect(publicSinkEntries).toHaveLength(0); + }); }); // ─── Usage Reporting ────────────────────────────────────────────────────── diff --git a/agentic/agentic-server/src/index.ts b/agentic/agentic-server/src/index.ts index a213f8ee59..6b4097a5bc 100644 --- a/agentic/agentic-server/src/index.ts +++ b/agentic/agentic-server/src/index.ts @@ -28,6 +28,7 @@ export type { AgenticServerStartOptions } from './server'; export { createAgenticServer } from './server'; export type { AgenticServerOptions, + InferenceAttribution, InferenceEntry, InferenceSink, ProviderConfig, diff --git a/agentic/agentic-server/src/router.ts b/agentic/agentic-server/src/router.ts index 9eeb059b77..03578994b5 100644 --- a/agentic/agentic-server/src/router.ts +++ b/agentic/agentic-server/src/router.ts @@ -21,7 +21,7 @@ import { transformEmbedRequest, transformEmbedResponse } from './transforms'; -import type { AgenticServerOptions, ResolvedProvider } from './types'; +import type { AgenticServerOptions, InferenceAttribution, ResolvedProvider } from './types'; const log = new Logger('agentic-server'); @@ -62,6 +62,26 @@ function upstreamErrorMessage( : `${provider.type} provider error ${status}`; } +/** + * Who the call is for and which task it is a cost of. Identity comes from the + * caller's identity headers; task linkage from the correlation headers the + * runtimes forward (`X-Invocation-Id`, `X-Job-Id`, `X-Attempt`, `X-Run-Id`). + * `X-Attempt` must be a non-negative integer or it is dropped. + */ +export const readAttribution = (req: any): InferenceAttribution => { + const rawAttempt = req.get('X-Attempt'); + const attempt = typeof rawAttempt === 'string' && /^\d+$/.test(rawAttempt) ? Number(rawAttempt) : undefined; + return { + databaseId: req.get('X-Database-Id'), + entityId: req.get('X-Entity-Id'), + actorId: req.get('X-Actor-Id'), + invocationId: req.get('X-Invocation-Id'), + jobId: req.get('X-Job-Id'), + attempt, + runId: req.get('X-Run-Id') + }; +}; + export const createRouter = (options: AgenticServerOptions): Router => { const router = Router(); // Metering is backend-agnostic: the caller injects an InferenceSink. When @@ -75,8 +95,8 @@ export const createRouter = (options: AgenticServerOptions): Router => { res.status(400).json({ error: { message: 'X-Database-Id is required' } }); return; } - const entityId = req.get('X-Entity-Id'); - const actorId = req.get('X-Actor-Id'); + const attribution = readAttribution(req); + const { entityId } = attribution; const requestProvider = req.get('X-LLM-Provider'); const startTime = process.hrtime.bigint(); @@ -134,7 +154,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || body.model || ''), provider: provider.type, service: 'chat', @@ -167,7 +187,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || body.model || ''), provider: provider.type, service: 'chat', @@ -196,7 +216,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || body.model || ''), provider: provider.type, service: 'chat', @@ -217,7 +237,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || ''), provider: provider.type, service: 'chat', @@ -249,8 +269,8 @@ export const createRouter = (options: AgenticServerOptions): Router => { res.status(400).json({ error: { message: 'X-Database-Id is required' } }); return; } - const entityId = req.get('X-Entity-Id'); - const actorId = req.get('X-Actor-Id'); + const attribution = readAttribution(req); + const { entityId } = attribution; const requestProvider = req.get('X-LLM-Provider'); const startTime = process.hrtime.bigint(); @@ -277,7 +297,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || body.model || ''), provider: provider.type, service: 'embed', @@ -302,7 +322,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || body.model || ''), provider: provider.type, service: 'embed', @@ -323,7 +343,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, entityId, actorId, + ...attribution, model: String(req.body?.model || ''), provider: provider.type, service: 'embed', @@ -348,8 +368,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { res.status(400).json({ error: { message: 'X-Database-Id is required' } }); return; } - const entityId = req.get('X-Entity-Id'); - const actorId = req.get('X-Actor-Id'); + const attribution = readAttribution(req); const { model, @@ -372,9 +391,7 @@ export const createRouter = (options: AgenticServerOptions): Router => { if (sink) { sink.logInference({ - databaseId, - entityId, - actorId, + ...attribution, model: String(model), provider: String(reportedProvider || 'unknown'), service: (service === 'embed' ? 'embed' : 'chat') as 'chat' | 'embed', diff --git a/agentic/agentic-server/src/server.ts b/agentic/agentic-server/src/server.ts index 165327b168..df3ca28797 100644 --- a/agentic/agentic-server/src/server.ts +++ b/agentic/agentic-server/src/server.ts @@ -7,11 +7,18 @@ export interface AgenticServerStartOptions extends AgenticServerOptions { port?: number; } -/** Identity headers trusted only when isPublic === false (private network). */ +/** Identity + task-correlation headers trusted only when isPublic === false + * (private network). The correlation headers say which invocation/job/attempt/ + * run a model call is a cost of; an external client must not be able to pin + * its usage onto someone else's task. */ const IDENTITY_HEADERS = [ 'x-database-id', 'x-entity-id', - 'x-actor-id' + 'x-actor-id', + 'x-invocation-id', + 'x-job-id', + 'x-attempt', + 'x-run-id' ] as const; /** diff --git a/agentic/agentic-server/src/types.ts b/agentic/agentic-server/src/types.ts index a4c0f2783c..f2f4553af7 100644 --- a/agentic/agentic-server/src/types.ts +++ b/agentic/agentic-server/src/types.ts @@ -24,6 +24,8 @@ export interface AgenticServerOptions { * When false (default), the server is deployed behind a private network * (Docker network, K8s service mesh) and trusts identity headers directly: * X-Database-Id, X-Entity-Id, X-Actor-Id + * and the task-correlation headers + * X-Invocation-Id, X-Job-Id, X-Attempt, X-Run-Id * When true, identity headers are stripped from incoming requests * (external clients cannot set tenant context). */ isPublic?: boolean; @@ -51,10 +53,24 @@ export interface InferenceSink { logInference(entry: InferenceEntry): void; } +/** The identity + task-linkage slice of an InferenceEntry, read from headers. */ +export type InferenceAttribution = Pick< + InferenceEntry, + 'databaseId' | 'entityId' | 'actorId' | 'invocationId' | 'jobId' | 'attempt' | 'runId' +>; + export interface InferenceEntry { databaseId: string; entityId?: string; actorId?: string; + /** Task linkage — the invocation / job / attempt / agent run this call is a + * cost of. Inference is COGS attributed to the task the customer was + * charged for; it is never a customer meter of its own. Absent when the + * call was served outside an invocation. */ + invocationId?: string; + jobId?: string; + attempt?: number; + runId?: string; model: string; provider: string; service: 'chat' | 'embed'; diff --git a/agentic/metering/__tests__/gateway.test.ts b/agentic/metering/__tests__/gateway.test.ts index 8ea24f92dc..6e00f1b307 100644 --- a/agentic/metering/__tests__/gateway.test.ts +++ b/agentic/metering/__tests__/gateway.test.ts @@ -1,13 +1,18 @@ import { ACTOR_ID_HEADER, + ATTEMPT_HEADER, buildIdentityHeaders, completionsBaseUrl, DATABASE_ID_HEADER, ENTITY_ID_HEADER, GATEWAY_API, + INVOCATION_ID_HEADER, + JOB_ID_HEADER, normalizeGatewayUrl, resolveMeteredGateway, - resolveMeteredModel} from '../src'; + resolveMeteredModel, + RUN_ID_HEADER +} from '../src'; const models = [{ id: 'anthropic/claude-sonnet-4', contextWindow: 200000, maxTokens: 8192 }]; @@ -37,6 +42,33 @@ describe('buildIdentityHeaders', () => { it('rejects a missing databaseId up front', () => { expect(() => buildIdentityHeaders({ databaseId: ' ' })).toThrow(/databaseId is required/); }); + + it('carries the task the run is a cost of: invocation, job, attempt, run', () => { + expect( + buildIdentityHeaders({ + databaseId: 'db-1', + invocationId: 'inv-1', + jobId: '4242', + attempt: 0, + runId: 'run-1' + }) + ).toEqual({ + [DATABASE_ID_HEADER]: 'db-1', + [INVOCATION_ID_HEADER]: 'inv-1', + [JOB_ID_HEADER]: '4242', + [ATTEMPT_HEADER]: '0', + [RUN_ID_HEADER]: 'run-1' + }); + }); + + it('refuses malformed task linkage rather than sending something the gateway drops', () => { + expect(() => buildIdentityHeaders({ databaseId: 'db-1', jobId: 'job-1' })).toThrow(/jobId must be decimal digits/); + expect(() => buildIdentityHeaders({ databaseId: 'db-1', attempt: -1 })).toThrow(/attempt must be a non-negative integer/); + expect(() => buildIdentityHeaders({ databaseId: 'db-1', attempt: 1.5 })).toThrow(/attempt must be a non-negative integer/); + expect(buildIdentityHeaders({ databaseId: 'db-1', invocationId: ' ', jobId: '', runId: '' })).toEqual({ + [DATABASE_ID_HEADER]: 'db-1' + }); + }); }); describe('normalizeGatewayUrl', () => { diff --git a/agentic/metering/src/identity.ts b/agentic/metering/src/identity.ts index 0aab508faf..399c49f533 100644 --- a/agentic/metering/src/identity.ts +++ b/agentic/metering/src/identity.ts @@ -1,10 +1,12 @@ /** - * Who a metered request belongs to. + * Who a metered request belongs to, and which task it is a cost of. * - * These are exactly the headers `agentic-server` reads (`X-Database-Id`, - * `X-Entity-Id`, `X-Actor-Id`) — the same identity lane the platform's own - * clients use — so a pi session lands in `inference_log` beside every other - * metered call rather than in a parallel accounting scheme. + * These are exactly the headers `agentic-server` reads — identity + * (`X-Database-Id`, `X-Entity-Id`, `X-Actor-Id`) plus task correlation + * (`X-Invocation-Id`, `X-Job-Id`, `X-Attempt`, `X-Run-Id`) — the same lane the + * platform's own clients use, so a pi session lands in `inference_log` beside + * every other metered call rather than in a parallel accounting scheme, and its + * tokens join the invocation the customer was actually charged for. * * Headers are only trustworthy where the gateway is not reachable by the * untrusted party: in-cluster, or behind an authenticated ingress that pins the @@ -19,6 +21,14 @@ export interface MeteredIdentity { entityId?: string; /** The actor on whose behalf the run executes. */ actorId?: string; + /** Invocation the run is executing, when the platform dispatched it. */ + invocationId?: string; + /** Queue job id for that invocation; decimal digits only. */ + jobId?: string; + /** Attempt number of that job; a non-negative integer. */ + attempt?: number; + /** The agent run itself. */ + runId?: string; /** * Bearer token for the gateway's ingress — run-scoped, not an account token. * Sent as `Authorization: Bearer `. @@ -29,6 +39,10 @@ export interface MeteredIdentity { export const DATABASE_ID_HEADER = 'X-Database-Id'; export const ENTITY_ID_HEADER = 'X-Entity-Id'; export const ACTOR_ID_HEADER = 'X-Actor-Id'; +export const INVOCATION_ID_HEADER = 'X-Invocation-Id'; +export const JOB_ID_HEADER = 'X-Job-Id'; +export const ATTEMPT_HEADER = 'X-Attempt'; +export const RUN_ID_HEADER = 'X-Run-Id'; /** * Build the identity headers for a metered request. @@ -46,6 +60,21 @@ export function buildIdentityHeaders(identity: MeteredIdentity): Record { expect(run.lanes.meteredModel!.selectedModel).toBe('gpt-5'); }); + it('books metered calls against the run: X-Run-Id rides on the gateway headers', () => { + const run = composeRun({ + runId: 'run-42', + metering: { mode: 'gateway', gatewayUrl, identity: { ...identity, invocationId: 'inv-1', jobId: '7', attempt: 1 }, models } + }); + + expect(run.lanes.meteredModel!.config.headers).toMatchObject({ + 'X-Run-Id': 'run-42', + 'X-Invocation-Id': 'inv-1', + 'X-Job-Id': '7', + 'X-Attempt': '1', + 'X-Entity-Id': 'ent-1' + }); + }); + + it('lets a host pin the run id on the identity explicitly', () => { + const run = composeRun({ + runId: 'run-42', + metering: { mode: 'gateway', gatewayUrl, identity: { ...identity, runId: 'run-pinned' }, models } + }); + + expect(run.lanes.meteredModel!.config.headers['X-Run-Id']).toBe('run-pinned'); + }); + it('picks the self-report lane for a run on the host’s own provider key', () => { const run = composeRun({ runId: 'run-1', diff --git a/agentic/pi/src/embed/lanes.ts b/agentic/pi/src/embed/lanes.ts index 2ec712f9e7..657ad0a471 100644 --- a/agentic/pi/src/embed/lanes.ts +++ b/agentic/pi/src/embed/lanes.ts @@ -81,14 +81,22 @@ export function composeRun(options: ComposeRunOptions): ComposedRun { extensions.push(lanes.log.extension); } + // The run id rides on the metering identity so the gateway can book every + // model call against this run (`X-Run-Id`); a host may still pin it explicitly. const metering = options.metering; if (metering?.mode === 'gateway') { - const { mode: _mode, ...meteredOptions } = metering; - lanes.meteredModel = createMeteredModelExtension(meteredOptions); + const { mode: _mode, identity, ...meteredOptions } = metering; + lanes.meteredModel = createMeteredModelExtension({ + ...meteredOptions, + identity: { runId: options.runId, ...identity } + }); extensions.push(lanes.meteredModel.extension); } else if (metering?.mode === 'self-report') { - const { mode: _mode, ...reportOptions } = metering; - lanes.usageReport = createUsageReportExtension(reportOptions); + const { mode: _mode, identity, ...reportOptions } = metering; + lanes.usageReport = createUsageReportExtension({ + ...reportOptions, + identity: { runId: options.runId, ...identity } + }); extensions.push(lanes.usageReport.extension); }