diff --git a/packages/adapters/grok-local/src/server/execute.test.ts b/packages/adapters/grok-local/src/server/execute.test.ts index b520bf13e6..3d7f5d6c9b 100644 --- a/packages/adapters/grok-local/src/server/execute.test.ts +++ b/packages/adapters/grok-local/src/server/execute.test.ts @@ -138,6 +138,70 @@ describe("grok_local execute", () => { expect(logs.map((entry) => entry.chunk)).not.toEqual([]); }); + it("reports real per-run token usage, marks it as per_run, and only surfaces cost for API billing", async () => { + const root = await makeTempRoot(); + + runProcessMock.mockImplementation(async () => ({ + exitCode: 0, + signal: null, + timedOut: false, + stdout: [ + JSON.stringify({ type: "text", data: "done" }), + JSON.stringify({ + type: "end", + stopReason: "EndTurn", + sessionId: "sess-1", + requestId: "req-1", + usage: { input_tokens: 2384, output_tokens: 261, cache_read_input_tokens: 23040 }, + total_cost_usd: 0.013246, + }), + ].join("\n"), + stderr: "", + })); + + const baseCtx: AdapterExecutionContext = { + runId: "run-1", + agent: { + id: "agent-1", + companyId: "company-1", + name: "Grok Agent", + adapterType: "grok_local", + adapterConfig: {}, + }, + runtime: { sessionId: null, sessionParams: null, sessionDisplayId: null, taskKey: null }, + config: { cwd: root }, + context: {}, + authToken: "run-token", + onLog: async () => {}, + }; + + // Subscription billing (no XAI_API_KEY): token usage is populated, but + // there is no marginal dollar cost so costUsd stays null. + const subscriptionResult = await execute(baseCtx); + expect(subscriptionResult).toMatchObject({ + usage: { inputTokens: 2384, outputTokens: 261, cachedInputTokens: 23040 }, + usageBasis: "per_run", + billingType: "subscription", + costUsd: null, + }); + + // API-key billing: same token usage, plus the real dollar cost. + const previousApiKey = process.env.XAI_API_KEY; + process.env.XAI_API_KEY = "test-key"; + try { + const apiResult = await execute(baseCtx); + expect(apiResult).toMatchObject({ + usage: { inputTokens: 2384, outputTokens: 261, cachedInputTokens: 23040 }, + usageBasis: "per_run", + billingType: "api", + costUsd: 0.013246, + }); + } finally { + if (previousApiKey === undefined) delete process.env.XAI_API_KEY; + else process.env.XAI_API_KEY = previousApiKey; + } + }); + it("cleans up staged assets when setup fails before the Grok process starts", async () => { const root = await makeTempRoot(); const instructionsPath = path.join(root, "managed", "AGENTS.md"); diff --git a/packages/adapters/grok-local/src/server/execute.ts b/packages/adapters/grok-local/src/server/execute.ts index 45ae407bc7..e82ff1cec6 100644 --- a/packages/adapters/grok-local/src/server/execute.ts +++ b/packages/adapters/grok-local/src/server/execute.ts @@ -541,10 +541,14 @@ export async function execute(ctx: AdapterExecutionContext): Promise { errorMessage: null, stopReason: "EndTurn", requestId: "req-1", + inputTokens: 0, + outputTokens: 0, + cachedInputTokens: 0, + costUsd: null, }); }); + it("extracts token usage and cost from the end event", () => { + const parsed = parseGrokJsonl([ + JSON.stringify({ type: "text", data: "hi" }), + JSON.stringify({ + type: "end", + stopReason: "EndTurn", + sessionId: "sess-1", + requestId: "req-1", + usage: { input_tokens: 21560, output_tokens: 960, cache_read_input_tokens: 25216 }, + total_cost_usd: 0.0564448, + }), + ].join("\n")); + + expect(parsed.inputTokens).toBe(21560); + expect(parsed.outputTokens).toBe(960); + expect(parsed.cachedInputTokens).toBe(25216); + expect(parsed.costUsd).toBe(0.0564448); + }); + it("reads structured error payloads", () => { const parsed = parseGrokJsonl([ JSON.stringify({ type: "error", error: { message: "Authentication required" } }), diff --git a/packages/adapters/grok-local/src/server/parse.ts b/packages/adapters/grok-local/src/server/parse.ts index 7fc415a239..5a36de0509 100644 --- a/packages/adapters/grok-local/src/server/parse.ts +++ b/packages/adapters/grok-local/src/server/parse.ts @@ -1,4 +1,4 @@ -import { asString, parseJson, parseObject } from "@paperclipai/adapter-utils/server-utils"; +import { asNumber, asString, parseJson, parseObject } from "@paperclipai/adapter-utils/server-utils"; import { applyTurnBoundary, createTurnBoundaryState } from "../shared/turn-boundary.js"; export interface ParsedGrokJsonl { @@ -8,6 +8,10 @@ export interface ParsedGrokJsonl { errorMessage: string | null; stopReason: string | null; requestId: string | null; + inputTokens: number; + outputTokens: number; + cachedInputTokens: number; + costUsd: number | null; } function errorText(value: unknown): string { @@ -31,6 +35,10 @@ export function parseGrokJsonl(stdout: string): ParsedGrokJsonl { let stopReason: string | null = null; let requestId: string | null = null; let errorMessage: string | null = null; + let inputTokens = 0; + let outputTokens = 0; + let cachedInputTokens = 0; + let costUsd: number | null = null; const thoughtParts: string[] = []; const textParts: string[] = []; const thoughtBoundary = createTurnBoundaryState(); @@ -59,6 +67,14 @@ export function parseGrokJsonl(stdout: string): ParsedGrokJsonl { sessionId = asString(event.sessionId, "").trim() || sessionId; stopReason = asString(event.stopReason, "").trim() || stopReason; requestId = asString(event.requestId, "").trim() || requestId; + const usage = parseObject(event.usage); + inputTokens = asNumber(usage.input_tokens, inputTokens); + outputTokens = asNumber(usage.output_tokens, outputTokens); + cachedInputTokens = asNumber(usage.cache_read_input_tokens, cachedInputTokens); + const totalCostUsd = event.total_cost_usd; + if (typeof totalCostUsd === "number" && Number.isFinite(totalCostUsd)) { + costUsd = totalCostUsd; + } continue; } @@ -75,6 +91,10 @@ export function parseGrokJsonl(stdout: string): ParsedGrokJsonl { errorMessage, stopReason, requestId, + inputTokens, + outputTokens, + cachedInputTokens, + costUsd, }; }