feat(agentic-server): carry invocation/job/attempt/run linkage into InferenceEntry

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 ──────────────────────────────────────────────────────

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,

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',

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;
 
 /**

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

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', () => {