Normalize xAI token usage before persistence

packages/fold-core/src/EventLog/Usage.ts

@@ -1,5 +1,5 @@
 import { Schema } from 'effect'
-import { Response } from 'effect/unstable/ai'
+import type { Response } from 'effect/unstable/ai'
 
 /** Best-effort token count reported by a model provider. Providers may omit any usage field. */
 export const UsageTokenCount = Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)).annotate({
@@ -41,12 +41,26 @@ export const UsageEncoded = Schema.Struct({
 }).annotate({ identifier: 'UsageEncoded' })
 export type UsageEncoded = typeof UsageEncoded.Type
 
-const encodeResponseUsage = Schema.encodeUnknownSync(Response.Usage)
 const decodeUsageEncoded = Schema.decodeUnknownSync(UsageEncoded)
 
+const nonNegativeTokenCount = (value: number | undefined): number | undefined =>
+	value !== undefined && Number.isFinite(value) && value >= 0 ? value : undefined
+
 /** Convert Effect AI usage into fold's tolerant durable usage shape. */
 export const usageFromResponseUsage = (usage: Response.Usage): UsageEncoded =>
-	decodeUsageEncoded(encodeResponseUsage(usage))
+	decodeUsageEncoded({
+		inputTokens: {
+			uncached: nonNegativeTokenCount(usage.inputTokens.uncached),
+			total: nonNegativeTokenCount(usage.inputTokens.total),
+			cacheRead: nonNegativeTokenCount(usage.inputTokens.cacheRead),
+			cacheWrite: nonNegativeTokenCount(usage.inputTokens.cacheWrite),
+		},
+		outputTokens: {
+			total: nonNegativeTokenCount(usage.outputTokens.total),
+			text: nonNegativeTokenCount(usage.outputTokens.text),
+			reasoning: nonNegativeTokenCount(usage.outputTokens.reasoning),
+		},
+	})
 
 /** Best estimate of total input tokens from whatever fields the provider reported. */
 export const usageInputTotal = (usage: UsageEncoded): number => {

packages/fold-core/test/EventLog/Usage.vi.test.ts

@@ -0,0 +1,36 @@
+import { describe, expect, it } from '@effect/vitest'
+import type { Response } from 'effect/unstable/ai'
+
+import { usageFromResponseUsage } from '../../src/EventLog/Usage'
+
+describe('usageFromResponseUsage', () => {
+	it('omits invalid derived token details while preserving provider totals', () => {
+		const usage: Response.Usage = {
+			inputTokens: {
+				uncached: 10,
+				total: 10,
+				cacheRead: 0,
+				cacheWrite: 0,
+			},
+			outputTokens: {
+				total: 8,
+				text: -4,
+				reasoning: 12,
+			},
+		}
+
+		expect(usageFromResponseUsage(usage)).toEqual({
+			inputTokens: {
+				uncached: 10,
+				total: 10,
+				cacheRead: 0,
+				cacheWrite: 0,
+			},
+			outputTokens: {
+				total: 8,
+				text: undefined,
+				reasoning: 12,
+			},
+		})
+	})
+})

packages/fold-xai/src/XaiModel.ts

@@ -1,9 +1,14 @@
 /** FoldModel factory for xAI's OpenAI-compatible inference API authenticated with OAuth. */
 import { OpenAiClient, OpenAiLanguageModel } from '@effect/ai-openai-compat'
+import type {
+	ChatCompletionChunk,
+	CreateResponse200,
+	CreateResponse200Sse,
+} from '@effect/ai-openai-compat/OpenAiClient'
 import * as NodeFileSystem from '@effect/platform-node/NodeFileSystem'
 import { customModel, resolveOpenAiReasoning } from '@humanlayer/fold-core'
 import type { FoldModel, ReasoningLevel } from '@humanlayer/fold-core'
-import { Context, Effect, Layer } from 'effect'
+import { Context, Effect, Layer, Option, Predicate, Schema, Stream } from 'effect'
 import type { Scope } from 'effect'
 import type { LanguageModel } from 'effect/unstable/ai'
 import { FetchHttpClient, HttpClient } from 'effect/unstable/http'
@@ -14,6 +19,48 @@ import { DEFAULT_XAI_MODEL_ID } from './XaiModelCatalog'
 
 export const XAI_API_URL = 'https://api.x.ai/v1'
 
+const TokenCount = Schema.Finite.check(Schema.isInt(), Schema.isGreaterThanOrEqualTo(0))
+const XaiCompletionTokenDetails = Schema.Struct({ reasoning_tokens: TokenCount })
+const decodeXaiCompletionTokenDetails = Schema.decodeUnknownOption(XaiCompletionTokenDetails)
+
+type XaiUsage = NonNullable<CreateResponse200['usage']>
+
+/**
+ * xAI reports `completion_tokens` as text-only while putting reasoning tokens in
+ * `completion_tokens_details`. OpenAI-compatible clients expect `completion_tokens` to include both.
+ */
+export const normalizeXaiChatCompletionUsage = (usage: XaiUsage): XaiUsage => {
+	const details = decodeXaiCompletionTokenDetails(usage.completion_tokens_details)
+	return Option.match(details, {
+		onNone: () => usage,
+		onSome: ({ reasoning_tokens: reasoningTokens }) => ({
+			...usage,
+			completion_tokens: usage.completion_tokens + reasoningTokens,
+		}),
+	})
+}
+
+const normalizeXaiResponse = <Response extends CreateResponse200 | ChatCompletionChunk>(
+	response: Response,
+): Response => {
+	if (Predicate.isNullish(response.usage)) return response
+	return { ...response, usage: normalizeXaiChatCompletionUsage(response.usage) }
+}
+
+const normalizeXaiStreamResponse = (response: CreateResponse200Sse): CreateResponse200Sse =>
+	typeof response === 'string' || '_tag' in response ? response : normalizeXaiResponse(response)
+
+/** Normalize xAI's token semantics before the stock OpenAI-compatible model derives usage details. */
+export const decorateXaiClient = (inner: OpenAiClient.Service): OpenAiClient.Service => ({
+	...inner,
+	createResponse: (options) =>
+		inner.createResponse(options).pipe(Effect.map(([body, response]) => [normalizeXaiResponse(body), response])),
+	createResponseStream: (options) =>
+		inner
+			.createResponseStream(options)
+			.pipe(Effect.map(([response, stream]) => [response, stream.pipe(Stream.map(normalizeXaiStreamResponse))])),
+})
+
 export type XaiModelOptions = {
 	readonly model?: string
 	readonly reasoning?: ReasoningLevel
@@ -35,8 +82,9 @@ export const makeXaiLanguageModel = (
 		const clientContext = yield* Layer.build(OpenAiClient.layer({ apiUrl: options.apiUrl ?? XAI_API_URL })).pipe(
 			Effect.provideService(HttpClient.HttpClient, withXaiAuth(base, auth)),
 		)
+		const client = decorateXaiClient(Context.get(clientContext, OpenAiClient.OpenAiClient))
 		return yield* OpenAiLanguageModel.make({ model: options.model ?? DEFAULT_XAI_MODEL_ID }).pipe(
-			Effect.provideService(OpenAiClient.OpenAiClient, Context.get(clientContext, OpenAiClient.OpenAiClient)),
+			Effect.provideService(OpenAiClient.OpenAiClient, client),
 		)
 	}).pipe(Effect.provide(NodeFileSystem.layer))
 

packages/fold-xai/test/Xai.vi.test.ts

@@ -8,6 +8,7 @@ import {
 	buildXaiAuthorizeUrl,
 	DEFAULT_XAI_MODEL_ID,
 	makeXaiAuthStore,
+	normalizeXaiChatCompletionUsage,
 	XAI_FRONTIER_MODELS,
 	XAI_BROWSER_REDIRECT_URI,
 	XAI_CLIENT_ID,
@@ -50,6 +51,37 @@ describe('xAI OAuth', () => {
 })
 
 describe('xaiModel', () => {
+	it('normalizes xAI text-only completion tokens for the OpenAI-compatible adapter', () => {
+		expect(
+			normalizeXaiChatCompletionUsage({
+				prompt_tokens: 641,
+				completion_tokens: 1,
+				total_tokens: 889,
+				prompt_tokens_details: { cached_tokens: 512 },
+				completion_tokens_details: { reasoning_tokens: 247 },
+			}),
+		).toMatchObject({
+			prompt_tokens: 641,
+			completion_tokens: 248,
+			total_tokens: 889,
+			completion_tokens_details: { reasoning_tokens: 247 },
+		})
+	})
+
+	it('normalizes xAI usage independently of the aggregate total', () => {
+		const usage = {
+			prompt_tokens: 10,
+			completion_tokens: 8,
+			total_tokens: 18,
+			completion_tokens_details: { reasoning_tokens: 3 },
+		}
+
+		expect(normalizeXaiChatCompletionUsage(usage)).toEqual({
+			...usage,
+			completion_tokens: 11,
+		})
+	})
+
 	it('exports the supported frontier catalog and defaults to its newest model', () => {
 		expect(XAI_FRONTIER_MODELS).toEqual([
 			{ modelId: 'grok-4.5', label: 'Grok 4.5' },