From 4ff070d2c19188d3a3b08a12de55ab5847beb92b Mon Sep 17 00:00:00 2001 From: halfaipg Date: Mon, 5 Oct 2026 22:35:03 -0400 Subject: [PATCH 1/2] fix(billing): display confirmed Codebase inference receipts --- package-lock.json | 1 + package.json | 1 + src/agent/AGENTS.md | 1 + src/agent/agent.ts | 13 +- src/agent/codebase-billing.ts | 47 +++++++ src/agent/codebase-stream.test.ts | 165 ++++++++++++++++++++++ src/agent/codebase-stream.ts | 227 ++++++++++++++++++++++++++++++ src/agent/safe-stream.ts | 12 +- src/commands/builtins/cost.ts | 28 ++-- src/commands/builtins/session.ts | 5 +- src/glue/client.ts | 16 ++- src/headless/run.ts | 2 +- src/ui-pi/app.ts | 22 ++- 13 files changed, 514 insertions(+), 26 deletions(-) create mode 100644 src/agent/codebase-billing.ts create mode 100644 src/agent/codebase-stream.test.ts create mode 100644 src/agent/codebase-stream.ts diff --git a/package-lock.json b/package-lock.json index cd22f5a..bad84f1 100644 --- a/package-lock.json +++ b/package-lock.json @@ -19,6 +19,7 @@ "glob": "^13.0.1", "ignore": "^7.0.0", "ink": "^5.2.1", + "openai": "6.26.0", "proper-lockfile": "^4.1.2", "react": "^18.3.1", "typebox": "^1.1.24", diff --git a/package.json b/package.json index fef0f30..31c5adc 100644 --- a/package.json +++ b/package.json @@ -73,6 +73,7 @@ "glob": "^13.0.1", "ignore": "^7.0.0", "ink": "^5.2.1", + "openai": "6.26.0", "proper-lockfile": "^4.1.2", "react": "^18.3.1", "typebox": "^1.1.24", diff --git a/src/agent/AGENTS.md b/src/agent/AGENTS.md index 68bdcb7..fe5f462 100644 --- a/src/agent/AGENTS.md +++ b/src/agent/AGENTS.md @@ -10,6 +10,7 @@ - Rewind and checkpoint behavior must keep conversation and filesystem state consistent. - Tournaments must isolate contestant worktrees and avoid merging failed/unsafe outputs. - Model selection should respect user config, live overrides, local/cloud capability, and BYOK settings. +- Proxy spending uses confirmed server receipts, never dollar estimates. Missing receipts are unconfirmed. Share the billing ledger across main/helper/subagent requests and preserve the safe proxy stream wrapper. Totals cover the current agent lifetime only. ## Work Guidance diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 72dd2e5..a8d7286 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -30,11 +30,13 @@ import { buildTools } from "../tools/registry.js"; import { TaskStore } from "../tools/task-store.js"; import type { ToolContext } from "../tools/types.js"; import { UserQueryStore } from "../user-queries/store.js"; +import { BillingLedger } from "./codebase-billing.js"; +import { createCodebaseStream } from "./codebase-stream.js"; import { type ResolvedConfig, resolveConfig } from "./config.js"; import { type Effort, resolveEffort } from "./effort.js"; import { fetchAvailableModels, type ModelOption } from "./model-list.js"; import { buildProjectFilesAddendum } from "./project-files.js"; -import { streamProxySafely } from "./safe-stream.js"; +import { createSafeProxyStream } from "./safe-stream.js"; import { buildSystemPrompt } from "./system-prompt.js"; const WRITE_TOOL_NAMES: ReadonlySet = new Set(["write_file", "edit_file", "multi_edit", "notebook_edit"]); @@ -111,6 +113,7 @@ export interface CreateAgentOptions { } export interface AgentBundle { + billing?: BillingLedger; agent: Agent; model: Model; source: ResolvedConfig["source"]; @@ -203,6 +206,8 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { const tokenManager = source === "proxy" ? new TokenManager({ store: credentials, oauthConfig: defaultOAuthConfig() }) : null; const getApiKey = tokenManager ? () => tokenManager.getAccessToken() : () => apiKey; + const billing = source === "proxy" ? new BillingLedger() : undefined; + const streamFn = billing ? createSafeProxyStream(createCodebaseStream(billing, model.baseUrl!)) : undefined; const config = persistedConfig; const permissions = new PermissionStore({ @@ -225,6 +230,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { fastModel: glueModels.fast, smartModel: glueModels.smart, getApiKey, + streamFn, }); // Pass model.contextWindow explicitly so proxy-synthesized models // (Codebase Auto, custom in-house IDs) get the real window instead of @@ -367,7 +373,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { ...(thinkingLevel && { thinkingLevel: thinkingLevel as Effort }), }, getApiKey, - ...(source === "proxy" && { streamFn: streamProxySafely }), + streamFn, beforeToolCall: (ctx, signal) => guardToolCall(ctx.toolCall.name, ctx.args, signal), afterToolCall: async (ctx, signal) => { await dispatchPostToolHooks( @@ -416,7 +422,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { ...(effort && { thinkingLevel: effort }), }, getApiKey, - ...(source === "proxy" && { streamFn: streamProxySafely }), + streamFn, transformContext: async (messages, signal) => { if (!compaction.needsCompaction(messages)) return withRelevantMemoryReminder(memory, messages); @@ -626,6 +632,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { void agentRef; return { agent, + billing, model, source, toolContext, diff --git a/src/agent/codebase-billing.ts b/src/agent/codebase-billing.ts new file mode 100644 index 0000000..e42b5a1 --- /dev/null +++ b/src/agent/codebase-billing.ts @@ -0,0 +1,47 @@ +export interface BillingReceipt { + creditsCharged: number | null; + usageEstimated: boolean; + model: string; +} + +export interface BillingSummary { + creditsCharged: number; + requests: number; + unconfirmed: number; + estimated: number; + models: string[]; +} + +export function parseBillingReceipt(value: unknown, model: string): BillingReceipt | undefined { + if (!value || typeof value !== "object") return undefined; + const row = value as Record; + if (typeof row.usage_estimated !== "boolean") return undefined; + const credits = row.credits_charged; + if (credits !== null && (typeof credits !== "number" || !Number.isSafeInteger(credits) || credits < 0)) { + return undefined; + } + return { creditsCharged: credits as number | null, usageEstimated: row.usage_estimated, model }; +} + +/** One entry per HTTP request, including sidecar and subagent calls. */ +export class BillingLedger { + private readonly entries: BillingReceipt[] = []; + record(receipt: BillingReceipt): void { + this.entries.push(receipt); + } + snapshot(): BillingSummary { + return { + creditsCharged: this.entries.reduce((sum, row) => sum + (row.creditsCharged ?? 0), 0), + requests: this.entries.length, + unconfirmed: this.entries.filter((row) => row.creditsCharged === null).length, + estimated: this.entries.filter((row) => row.usageEstimated).length, + models: [...new Set(this.entries.map((row) => row.model))], + }; + } +} + +export function formatBillingSummary(summary: BillingSummary): string { + if (summary.requests === 0) return "No confirmed charges yet"; + const unknown = summary.unconfirmed ? ` + ${summary.unconfirmed} unconfirmed` : ""; + return `${summary.creditsCharged.toLocaleString()} credits charged${unknown}`; +} diff --git a/src/agent/codebase-stream.test.ts b/src/agent/codebase-stream.test.ts new file mode 100644 index 0000000..ce95fd7 --- /dev/null +++ b/src/agent/codebase-stream.test.ts @@ -0,0 +1,165 @@ +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import type { Model } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { BillingLedger, formatBillingSummary, parseBillingReceipt } from "./codebase-billing.js"; +import { createCodebaseStream } from "./codebase-stream.js"; + +let server: Server | undefined; +afterEach(async () => { + server?.closeAllConnections(); + if (server) await new Promise((resolve) => server!.close(() => resolve())); +}); +async function serve(chunks: unknown[], status = 200) { + let request: Record | undefined; + server = createServer(async (req, res) => { + let raw = ""; + for await (const bytes of req) raw += bytes; + request = JSON.parse(raw); + res.writeHead(status, { "content-type": status === 200 ? "text/event-stream" : "application/json" }); + if (status !== 200) return res.end(JSON.stringify(chunks[0])); + for (const chunk of chunks) res.write(`data: ${JSON.stringify(chunk)}\n\n`); + res.end("data: [DONE]\n\n"); + }); + await new Promise((resolve) => server!.listen(0, "127.0.0.1", resolve)); + const model: Model = { + id: "d4f", + name: "Auto", + provider: "codebase", + api: "openai-completions", + baseUrl: `http://127.0.0.1:${(server.address() as AddressInfo).port}`, + reasoning: false, + input: ["text"], + contextWindow: 131072, + maxTokens: 8192, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + }; + const ledger = new BillingLedger(); + const stream = createCodebaseStream(ledger, model.baseUrl!); + return { + ledger, + model, + stream, + request: () => request, + run: (signal?: AbortSignal) => + stream( + model, + { systemPrompt: "Be concise", messages: [{ role: "user", content: "hi", timestamp: 1 }] }, + { apiKey: "fake-key", signal }, + ), + }; +} +const final = { + model: "deepseek-flash", + choices: [], + usage: { + prompt_tokens: 100, + completion_tokens: 10, + total_tokens: 110, + prompt_tokens_details: { cached_tokens: 80, cache_write_tokens: 5 }, + }, + codebase_billing: { credits_charged: 3, usage_estimated: false }, +}; + +describe("Codebase receipt transport", () => { + it("uses the SDK stream with terminal usage-only receipts and the actual fallback model", async () => { + const h = await serve([ + { model: "deepseek-flash", choices: [{ delta: { content: "Hello" }, finish_reason: "stop" }] }, + final, + final, + ]); + const result = await h.run().result(); + expect(result.content).toEqual([{ type: "text", text: "Hello" }]); + expect(result.responseModel).toBe("deepseek-flash"); + expect(result.usage).toMatchObject({ input: 15, output: 10, cacheRead: 80, cacheWrite: 5 }); + expect(h.ledger.snapshot()).toEqual({ + creditsCharged: 3, + requests: 1, + estimated: 0, + unconfirmed: 0, + models: ["deepseek-flash"], + }); + expect(h.request()).toMatchObject({ stream: true, stream_options: { include_usage: true } }); + }); + it("never calls a missing or null receipt free", async () => { + const h = await serve([{ choices: [{ delta: { content: "Hello" }, finish_reason: "stop" }] }]); + await h.run().result(); + expect(formatBillingSummary(h.ledger.snapshot())).toBe("0 credits charged + 1 unconfirmed"); + expect(parseBillingReceipt({ credits_charged: null, usage_estimated: true }, "d4f")?.creditsCharged).toBeNull(); + }); + it("keeps the receipt sent before an interrupted-stream SDK error", async () => { + const h = await serve([final, { error: { type: "provider_error", message: "upstream unavailable" } }]); + expect((await h.run().result()).stopReason).toBe("error"); + expect(h.ledger.snapshot()).toMatchObject({ creditsCharged: 3, requests: 1, unconfirmed: 0 }); + }); + it("shows an admission denial without a retry or invented debit", async () => { + const h = await serve([{ error: "insufficient_credits", error_description: "Not enough credits" }], 402); + expect((await h.run().result()).stopReason).toBe("error"); + expect(h.ledger.snapshot()).toMatchObject({ creditsCharged: 0, requests: 1, unconfirmed: 1 }); + }); + it("keeps cancellation unconfirmed rather than zero-cost", async () => { + const h = await serve([]); + const controller = new AbortController(); + controller.abort(); + expect((await h.run(controller.signal).result()).stopReason).toBe("aborted"); + expect(h.ledger.snapshot().unconfirmed).toBe(1); + }); + it("preserves SDK tool deltas", async () => { + const h = await serve([ + { + choices: [ + { + delta: { + tool_calls: [{ index: 0, id: "call1", function: { name: "write_file", arguments: '{"path":' } }], + }, + }, + ], + }, + { + choices: [ + { + delta: { tool_calls: [{ index: 0, function: { arguments: '"a.txt"}' } }] }, + finish_reason: "tool_calls", + }, + ], + }, + final, + ]); + const result = await h.run().result(); + expect(result.stopReason).toBe("toolUse"); + expect(result.content).toEqual([ + { type: "toolCall", id: "call1", name: "write_file", arguments: { path: "a.txt" } }, + ]); + }); + it("rejects incomplete tool JSON without a completed tool event", async () => { + const h = await serve([ + { + choices: [ + { + delta: { + tool_calls: [{ index: 0, id: "call1", function: { name: "write_file", arguments: '{"path":' } }], + }, + finish_reason: "length", + }, + ], + }, + final, + ]); + const events = []; + for await (const event of h.run()) events.push(event); + expect(events.some((event) => event.type === "toolcall_end")).toBe(false); + expect(events.at(-1)?.type).toBe("error"); + expect(h.ledger.snapshot().requests).toBe(1); + }); + it("never sends a Codebase key to a helper endpoint override", async () => { + const h = await serve([]); + const result = await h + .stream({ ...h.model, baseUrl: "https://example.com/v1" }, { messages: [] }, { apiKey: "fake-key" }) + .result(); + expect(result.errorMessage).toContain("different endpoint"); + expect(h.request()).toBeUndefined(); + }); + it.each([-1, 0.5, Number.NaN, "3", undefined])("rejects invalid receipt amount %s", (value) => { + expect(parseBillingReceipt({ credits_charged: value, usage_estimated: false }, "d4f")).toBeUndefined(); + }); +}); diff --git a/src/agent/codebase-stream.ts b/src/agent/codebase-stream.ts new file mode 100644 index 0000000..68d27bd --- /dev/null +++ b/src/agent/codebase-stream.ts @@ -0,0 +1,227 @@ +import { + type AssistantMessage, + createAssistantMessageEventStream, + type Model, + parseStreamingJson, + type SimpleStreamOptions, + type StreamFunction, + type ToolCall, +} from "@earendil-works/pi-ai"; +import { convertMessages } from "@earendil-works/pi-ai/openai-completions"; +import OpenAI from "openai"; +import type { ChatCompletionChunk, ChatCompletionCreateParamsStreaming } from "openai/resources/chat/completions.js"; +import { type BillingLedger, type BillingReceipt, parseBillingReceipt } from "./codebase-billing.js"; + +// Codebase normalizes every upstream to Chat Completions. Keep pi's message +// conversion and the official SDK's SSE parser; only the billing extension is ours. +const COMPAT = { + supportsStore: false, + supportsDeveloperRole: false, + supportsReasoningEffort: true, + supportsUsageInStreaming: true, + maxTokensField: "max_tokens" as const, + requiresToolResultName: true, + requiresAssistantAfterToolResult: false, + requiresThinkingAsText: false, + requiresReasoningContentOnAssistantMessages: false, + thinkingFormat: "openai" as const, + openRouterRouting: {}, + vercelGatewayRouting: {}, + zaiToolStream: false, + supportsStrictMode: false, + sendSessionAffinityHeaders: false, + supportsLongCacheRetention: false, +}; + +type BillingChunk = ChatCompletionChunk & { codebase_billing?: unknown }; +type BilledMessage = AssistantMessage & { codebaseBilling?: BillingReceipt }; + +export function createCodebaseStream( + ledger: BillingLedger, + baseUrl: string, +): StreamFunction { + return (model, context, options = {}) => { + const stream = createAssistantMessageEventStream(); + const output: BilledMessage = { + role: "assistant", + api: "openai-completions", + provider: model.provider, + model: model.id, + content: [], + stopReason: "stop", + timestamp: Date.now(), + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; + void (async () => { + let receipt: BillingReceipt | undefined; + let recorded = false; + const tools = new Map(); + try { + if (!model.baseUrl || new URL(model.baseUrl).href !== new URL(baseUrl).href) + throw new Error("Codebase credentials cannot be used with a different endpoint."); + const client = new OpenAI({ + apiKey: options.apiKey, + baseURL: model.baseUrl, + maxRetries: 0, + timeout: options.timeoutMs ?? 600_000, + defaultHeaders: { ...model.headers, ...options.headers }, + }); + const chatModel = { ...model, api: "openai-completions", reasoning: false } as Model<"openai-completions">; + let payload: ChatCompletionCreateParamsStreaming = { + model: model.id, + messages: convertMessages(chatModel, context, COMPAT), + stream: true, + stream_options: { include_usage: true }, + max_tokens: Math.min(options.maxTokens ?? model.maxTokens, model.maxTokens), + ...(options.temperature !== undefined ? { temperature: options.temperature } : {}), + ...(context.tools?.length + ? { + tools: context.tools.map((tool) => ({ + type: "function" as const, + function: { + name: tool.name, + description: tool.description, + parameters: tool.parameters as unknown as Record, + }, + })), + } + : {}), + }; + const replacement = await options.onPayload?.(payload, model); + if (replacement !== undefined) payload = replacement as ChatCompletionCreateParamsStreaming; + const { data, response } = await client.chat.completions + .create(payload, { signal: options.signal }) + .withResponse(); + await options.onResponse?.( + { status: response.status, headers: Object.fromEntries(response.headers) }, + model, + ); + stream.push({ type: "start", partial: output }); + for await (const raw of data) { + const chunk = raw as BillingChunk; + if (chunk.model) output.responseModel = chunk.model; + output.responseId ||= chunk.id; + const parsed = parseBillingReceipt( + chunk.codebase_billing, + chunk.model || output.responseModel || model.id, + ); + if (parsed) receipt = parsed; + if (chunk.usage) { + const cached = chunk.usage.prompt_tokens_details?.cached_tokens ?? 0; + const writes = + (chunk.usage.prompt_tokens_details as { cache_write_tokens?: number } | undefined) + ?.cache_write_tokens ?? 0; + output.usage.input = Math.max(0, chunk.usage.prompt_tokens - cached - writes); + output.usage.output = chunk.usage.completion_tokens; + output.usage.cacheRead = cached; + output.usage.cacheWrite = writes; + output.usage.totalTokens = chunk.usage.total_tokens; + } + const choice = chunk.choices?.[0]; + if (!choice) continue; + if (choice.finish_reason === "tool_calls") output.stopReason = "toolUse"; + else if (choice.finish_reason === "length") output.stopReason = "length"; + else if (choice.finish_reason === "content_filter") + throw new Error("The provider refused this response."); + const delta = choice.delta; + if (delta.content) appendText("text", delta.content); + const reasoning = delta as typeof delta & { reasoning_content?: string; reasoning?: string }; + if (reasoning.reasoning_content || reasoning.reasoning) + appendText("thinking", reasoning.reasoning_content || reasoning.reasoning || ""); + for (const call of delta.tool_calls ?? []) { + let entry = tools.get(call.index); + if (!entry) { + entry = { + block: { type: "toolCall", id: call.id || "", name: call.function?.name || "", arguments: {} }, + args: "", + }; + tools.set(call.index, entry); + output.content.push(entry.block); + stream.push({ + type: "toolcall_start", + contentIndex: output.content.indexOf(entry.block), + partial: output, + }); + } + if (call.id) entry.block.id = call.id; + if (call.function?.name) entry.block.name = call.function.name; + entry.args += call.function?.arguments || ""; + entry.block.arguments = parseStreamingJson(entry.args); + stream.push({ + type: "toolcall_delta", + contentIndex: output.content.indexOf(entry.block), + delta: call.function?.arguments || "", + partial: output, + }); + } + } + if (options.signal?.aborted) throw new Error("Request cancelled"); + for (const [contentIndex, block] of output.content.entries()) { + if (block.type === "toolCall") { + const entry = [...tools.values()].find((candidate) => candidate.block === block); + block.arguments = JSON.parse(entry?.args || "{}"); + stream.push({ type: "toolcall_end", contentIndex, toolCall: block, partial: output }); + } else if (block.type === "text") + stream.push({ type: "text_end", contentIndex, content: block.text, partial: output }); + else if (block.type === "thinking") + stream.push({ type: "thinking_end", contentIndex, content: block.thinking, partial: output }); + } + finishBilling(); + stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output }); + } catch (error) { + const data = + error instanceof OpenAI.APIError ? (error.error as Record | undefined) : undefined; + receipt = parseBillingReceipt(data?.codebase_billing, output.responseModel || model.id) ?? receipt; + output.stopReason = options.signal?.aborted ? "aborted" : "error"; + output.errorMessage = options.signal?.aborted + ? "Request cancelled; check Usage for its final charge." + : error instanceof Error + ? error.message + : "Inference request failed"; + finishBilling(); + stream.push({ type: "error", reason: output.stopReason, error: output }); + } finally { + stream.end(); + } + + function finishBilling(): void { + if (recorded) return; + recorded = true; + output.codebaseBilling = receipt ?? { + creditsCharged: null, + usageEstimated: false, + model: output.responseModel || model.id, + }; + ledger.record(output.codebaseBilling); + } + function appendText(kind: "text" | "thinking", text: string): void { + let block = output.content.at(-1); + if (!block || block.type !== kind) { + block = kind === "text" ? { type: "text", text: "" } : { type: "thinking", thinking: "" }; + output.content.push(block); + stream.push({ + type: kind === "text" ? "text_start" : "thinking_start", + contentIndex: output.content.length - 1, + partial: output, + }); + } + if (block.type === "text") block.text += text; + else if (block.type === "thinking") block.thinking += text; + stream.push({ + type: kind === "text" ? "text_delta" : "thinking_delta", + contentIndex: output.content.length - 1, + delta: text, + partial: output, + }); + } + })(); + return stream; + }; +} diff --git a/src/agent/safe-stream.ts b/src/agent/safe-stream.ts index aeca6b2..299ad15 100644 --- a/src/agent/safe-stream.ts +++ b/src/agent/safe-stream.ts @@ -4,6 +4,8 @@ import { type AssistantMessageEventStream, createAssistantMessageEventStream, streamSimple, + type SimpleStreamOptions, + type StreamFunction, } from "@earendil-works/pi-ai"; const DSML_MARKER = "<|dsml|"; @@ -51,6 +53,14 @@ export function sanitizeAssistantMessage(message: AssistantMessage): AssistantMe * executor see the same safe response. */ export function streamProxySafely(...args: Parameters): ReturnType { + return createSafeProxyStream(streamSimple)(...args); +} + +export function createSafeProxyStream(upstreamFn: StreamFunction): typeof streamSimple { + return (...args) => streamWithSafety(upstreamFn, ...args); +} + +function streamWithSafety(upstreamFn: StreamFunction, ...args: Parameters): ReturnType { const [model, context, options] = args; const output = createAssistantMessageEventStream(); const controller = new AbortController(); @@ -78,7 +88,7 @@ export function streamProxySafely(...args: Parameters): Ret }, timeoutMs); timer.unref?.(); - const upstream = streamSimple(model, context, { ...options, signal: controller.signal }); + const upstream = upstreamFn(model, context, { ...options, signal: controller.signal }); void (async () => { try { for await (const event of upstream) { diff --git a/src/commands/builtins/cost.ts b/src/commands/builtins/cost.ts index 1ac7d84..ccdafc1 100644 --- a/src/commands/builtins/cost.ts +++ b/src/commands/builtins/cost.ts @@ -1,3 +1,4 @@ +import { formatBillingSummary } from "../../agent/codebase-billing.js"; import type { Command } from "../types.js"; export const cost: Command = { @@ -6,24 +7,27 @@ export const cost: Command = { handler: (_args, ctx) => { const { state, bundle } = ctx; const u = state.usage; - const turns = state.messages.filter((m) => m.role === "assistant").length; - const promptTokens = u.input + u.cacheRead; - const hitRate = promptTokens > 0 ? `${((u.cacheRead / promptTokens) * 100).toFixed(0)}%` : "—"; - const turnAvg = turns > 0 ? u.cost.total / turns : 0; - const proxyNote = bundle.source === "proxy" ? " (proxied via codebase.design)" : ""; - const usageUnavailable = - bundle.source === "proxy" && - u.input + u.output + u.cacheRead + u.cacheWrite + u.totalTokens === 0 && - turns > 0; - if (usageUnavailable) { + if (bundle.source === "proxy") { + const summary = bundle.billing?.snapshot(); ctx.emit( - "Session token and dollar usage is unavailable because the proxy did not return provider usage. Run /usage for account credits.", + [ + `Codebase spending: ${summary ? formatBillingSummary(summary) : "Charge information unavailable"}`, + "Since this CLI agent started, including its helper and subagent calls.", + ...(summary?.estimated ? [`${summary.estimated} request(s) used estimated usage.`] : []), + ...(summary?.models.length ? [`Served models: ${summary.models.join(", ")}`] : []), + `Tokens: ${u.input} input, ${u.output} output, ${u.cacheRead} cache read, ${u.cacheWrite} cache write`, + "Settings > Usage shows the authoritative account balance and final cancelled-request charges.", + ].join("\n"), ); return { handled: true }; } + const turns = state.messages.filter((m) => m.role === "assistant").length; + const promptTokens = u.input + u.cacheRead; + const hitRate = promptTokens > 0 ? `${((u.cacheRead / promptTokens) * 100).toFixed(0)}%` : "—"; + const turnAvg = turns > 0 ? u.cost.total / turns : 0; const lines = [ - `Session cost: $${u.cost.total.toFixed(4)}${proxyNote}`, + `Session cost estimate: $${u.cost.total.toFixed(4)}`, "", "Tokens:", ` Input ${padNum(u.input, 8)} ($${u.cost.input.toFixed(4)})`, diff --git a/src/commands/builtins/session.ts b/src/commands/builtins/session.ts index 6f152ad..1969e44 100644 --- a/src/commands/builtins/session.ts +++ b/src/commands/builtins/session.ts @@ -1,3 +1,4 @@ +import { formatBillingSummary } from "../../agent/codebase-billing.js"; import type { Command } from "../types.js"; export const clear: Command = { @@ -65,7 +66,9 @@ export const session: Command = { usageUnavailable ? "usage: unavailable from proxy" : `usage: ↓${u.input} ↑${u.output} cache ↓${u.cacheRead}/↑${u.cacheWrite}`, - usageUnavailable ? "cost: unavailable from proxy" : `cost: $${u.cost.total.toFixed(4)}`, + bundle.source === "proxy" + ? `spending: ${bundle.billing ? formatBillingSummary(bundle.billing.snapshot()) : "Charge information unavailable"} (current agent)` + : `cost: $${u.cost.total.toFixed(4)}`, `source: ${bundle.source}`, ]; ctx.emit(lines.join("\n")); diff --git a/src/glue/client.ts b/src/glue/client.ts index c3991a1..5e9711e 100644 --- a/src/glue/client.ts +++ b/src/glue/client.ts @@ -1,4 +1,11 @@ -import { completeSimple, getModel, type KnownProvider, type Model } from "@earendil-works/pi-ai"; +import { + completeSimple, + getModel, + type KnownProvider, + type Model, + type SimpleStreamOptions, + type StreamFunction, +} from "@earendil-works/pi-ai"; /** * Cheap-model sidecar. The agent's main model handles tool-using work; @@ -26,6 +33,7 @@ export interface GlueOptions { getApiKey: () => Promise | string; /** Default 8000 — glue calls cap their reply length so a runaway summary doesn't hang the UI. */ maxTokens?: number; + streamFn?: StreamFunction; } export class GlueClient { @@ -46,7 +54,11 @@ export class GlueClient { signal: AbortSignal | undefined, ): Promise { const apiKey = await this.options.getApiKey(); - const message = await completeSimple( + const complete = this.options.streamFn + ? (model: Model, context: Parameters[1], options: SimpleStreamOptions) => + this.options.streamFn!(model, context, options).result() + : completeSimple; + const message = await complete( model, { systemPrompt: system, diff --git a/src/headless/run.ts b/src/headless/run.ts index 7390bf0..922c31a 100644 --- a/src/headless/run.ts +++ b/src/headless/run.ts @@ -398,7 +398,7 @@ export async function runHeadless(opts: HeadlessOptions): Promise { receipt, receiptRecord, }); - out(`${JSON.stringify(payload)}\n`); + out(`${JSON.stringify({ ...payload, ...(bundle.billing ? { billing: bundle.billing.snapshot() } : {}) })}\n`); } return exitCode; diff --git a/src/ui-pi/app.ts b/src/ui-pi/app.ts index 3a743f5..df072fa 100644 --- a/src/ui-pi/app.ts +++ b/src/ui-pi/app.ts @@ -11,6 +11,7 @@ import { type TUI, } from "@earendil-works/pi-tui"; import { type AgentBundle, createAgent } from "../agent/agent.js"; +import { formatBillingSummary } from "../agent/codebase-billing.js"; import { CHARS_PER_TOKEN, estimateContextTokens, streamingChars } from "../agent/context-estimate.js"; import { listRewindPoints, type RewindPoint, truncateBefore } from "../agent/conversation-rewind.js"; import { fetchAvailableModels } from "../agent/model-list.js"; @@ -1304,7 +1305,12 @@ export class App extends Container { // compaction engine, which already trusts model.contextWindow. const contextWindow = this.bundle.model.contextWindow || 200_000; const ctxPct = contextWindow > 0 ? Math.min(100, Math.round((usedTokens / contextWindow) * 100)) : 0; - this.statusBar.setMetrics(ctxPct, this.usage.cost.total, this.computeTokRate()); + this.statusBar.setMetrics( + ctxPct, + this.usage.cost.total, + this.computeTokRate(), + this.bundle.billing ? formatBillingSummary(this.bundle.billing.snapshot()) : undefined, + ); this.contextWarning.setPercent(ctxPct); } @@ -1575,15 +1581,14 @@ class StatusBar extends Container { private currentStatus = "idle"; private ctxPercent = 0; private cost = 0; + private billingText: string | undefined; private tokRate: number | undefined; private throbberTick = 0; /** Active cycling verb shown while currentStatus is "thinking". */ private verb = THINKING_VERBS[0]; private verbTimer: NodeJS.Timeout | undefined; private readonly onTick: () => void; - /** Per-turn cost is only meaningful on metered (BYOK) sessions — the - * proxy bills a flat subscription and returns no usage, so $0.0000 there - * is noise, not information. Hidden for proxy. */ + /** BYOK shows token-cost estimates; Codebase shows server-confirmed credits. */ private readonly showCost: boolean; constructor(modelName: string, cwd: string, onTick: () => void = () => undefined, showCost = true) { super(); @@ -1632,9 +1637,10 @@ class StatusBar extends Container { this.line.setText(this.format()); this.line.invalidate(); } - setMetrics(ctxPercent: number, cost: number, tokRate?: number): void { + setMetrics(ctxPercent: number, cost: number, tokRate?: number, billingText?: string): void { this.ctxPercent = ctxPercent; this.cost = cost; + this.billingText = billingText; this.tokRate = tokRate; this.line.setText(this.format()); this.line.invalidate(); @@ -1674,7 +1680,11 @@ class StatusBar extends Container { const bar = ctxBar(this.ctxPercent); const ctxText = colorByThreshold(`${bar} ${this.ctxPercent}%`, this.ctxPercent); const tokPart = this.tokRate !== undefined ? ` · ${this.tokRate} tok/s` : ""; - const costPart = this.showCost ? ` · $${formatCost(this.cost)}` : ""; + const costPart = this.billingText + ? ` · ${this.billingText}` + : this.showCost + ? ` · $${formatCost(this.cost)}` + : ""; // Model name reads at default brightness as the anchor; the rest of // the meta recedes to dim so the line is glanceable, not a wall. const meta = ansi.dim(`${this.cwdLabel} · ctx ${ctxText}${tokPart}${costPart}`); From 22b87882ffdbd4af1c9e98028afd5501fa81b129 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Mon, 5 Oct 2026 22:52:03 -0400 Subject: [PATCH 2/2] style(billing): format the safe proxy stream wrapper --- src/agent/safe-stream.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/agent/safe-stream.ts b/src/agent/safe-stream.ts index 299ad15..90331c6 100644 --- a/src/agent/safe-stream.ts +++ b/src/agent/safe-stream.ts @@ -3,9 +3,9 @@ import { type AssistantMessageEvent, type AssistantMessageEventStream, createAssistantMessageEventStream, - streamSimple, type SimpleStreamOptions, type StreamFunction, + streamSimple, } from "@earendil-works/pi-ai"; const DSML_MARKER = "<|dsml|"; @@ -60,7 +60,10 @@ export function createSafeProxyStream(upstreamFn: StreamFunction streamWithSafety(upstreamFn, ...args); } -function streamWithSafety(upstreamFn: StreamFunction, ...args: Parameters): ReturnType { +function streamWithSafety( + upstreamFn: StreamFunction, + ...args: Parameters +): ReturnType { const [model, context, options] = args; const output = createAssistantMessageEventStream(); const controller = new AbortController();