diff --git a/internal/cloud/access-control.ts b/internal/cloud/access-control.ts index fd884469..6fa404f0 100644 --- a/internal/cloud/access-control.ts +++ b/internal/cloud/access-control.ts @@ -2149,11 +2149,26 @@ function normalizeAgentSessionTraceResult(value) { finalResponse: redactTraceFinalResponse(record.finalResponse), traceSummary: redactTraceSummary(record.traceSummary), agentRun: redactAgentRunSummary(agentRun), + userBillingReservation: redactBillingReservation(record.userBillingReservation), valuesRedacted: true, secretMaterialStored: false }); return result.agentRun?.runId || result.finalResponse || result.traceSummary ? result : null; } + +function redactBillingReservation(value) { + const reservation = normalizeObject(value); + const reservationId = textOr(reservation.reservationId, ""); + if (!reservationId) return null; + return pruneEmpty({ + reservationId, + estimatedCredits: numberOrNull(reservation.estimatedCredits), + estimatedTokens: numberOrNull(reservation.estimatedTokens), + expiresAt: textOr(reservation.expiresAt, ""), + valuesRedacted: true, + secretMaterialStored: false + }); +} function agentSessionTraceEvidence(session, traceId) { const id = safeTraceIdLocal(traceId); if (!id) return null; diff --git a/internal/cloud/code-agent-agentrun-adapter.ts b/internal/cloud/code-agent-agentrun-adapter.ts index cb1db5ca..68736500 100644 --- a/internal/cloud/code-agent-agentrun-adapter.ts +++ b/internal/cloud/code-agent-agentrun-adapter.ts @@ -694,7 +694,8 @@ export async function loadPersistedAgentRunResult(traceId, options = {}) { capabilityLevel: AGENTRUN_CAPABILITY_LEVEL, sessionMode: AGENTRUN_SESSION_MODE, implementationType: AGENTRUN_IMPLEMENTATION_TYPE, - agentRun: { ...agentRun, adapter: ADAPTER_ID, traceId: safeId, commandId: text(agentRun.commandId) || null, valuesPrinted: false }, + ...(agentRun?.userBillingReservation ? { userBillingReservation: agentRun.userBillingReservation } : {}), + agentRun: { ...withoutUserBillingReservation(agentRun), adapter: ADAPTER_ID, traceId: safeId, commandId: text(agentRun.commandId) || null, valuesPrinted: false }, valuesPrinted: false }; } @@ -702,8 +703,13 @@ export async function loadPersistedAgentRunResult(traceId, options = {}) { function agentRunSeedFromSession(session, traceId) { const snapshot = session?.session && typeof session.session === "object" ? session.session : null; const traceResults = snapshot?.traceResults && typeof snapshot.traceResults === "object" ? snapshot.traceResults : null; - const traceAgentRun = traceResults?.[traceId]?.agentRun && typeof traceResults[traceId].agentRun === "object" ? traceResults[traceId].agentRun : null; - if (traceAgentRun?.runId) return traceAgentRun; + const traceResult = traceResults?.[traceId] && typeof traceResults[traceId] === "object" ? traceResults[traceId] : null; + const traceAgentRun = traceResult?.agentRun && typeof traceResult.agentRun === "object" ? traceResult.agentRun : null; + if (traceAgentRun?.runId) { + return traceResult?.userBillingReservation + ? { ...traceAgentRun, userBillingReservation: traceResult.userBillingReservation } + : traceAgentRun; + } const topLevelAgentRun = snapshot?.agentRun && typeof snapshot.agentRun === "object" ? snapshot.agentRun : null; return topLevelAgentRun?.runId ? topLevelAgentRun : null; } @@ -750,6 +756,12 @@ export function agentRunSessionEvidence(payload = {}) { }; } +function withoutUserBillingReservation(value = {}) { + if (!value || typeof value !== "object") return value; + const { userBillingReservation, ...rest } = value; + return rest; +} + async function ensureAgentRunSessionPersistent({ fetchImpl, managerUrl, sessionId, env, traceId, backendProfile, traceStore }) { const defaultPolicy = firstNonEmpty(env.HWLAB_CODE_AGENT_AGENTRUN_SESSION_STORAGE, "persistent"); if (defaultPolicy !== "persistent") return; diff --git a/internal/cloud/server-agent-chat.test.ts b/internal/cloud/server-agent-chat.test.ts index a8495ab6..6bb80d33 100644 --- a/internal/cloud/server-agent-chat.test.ts +++ b/internal/cloud/server-agent-chat.test.ts @@ -92,6 +92,7 @@ test("AgentRun adapter filters resource tools and credentials through access cap HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${port}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_RUNTIME_API_URL: "http://hwlab-cloud-api.hwlab-v02.svc.cluster.local:6667", @@ -468,6 +469,7 @@ test("cloud api /v1/agent/chat delegates v0.2 turns to AgentRun v0.1 over adapte HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_CODE_AGENT_DEEPSEEK_MODEL: "deepseek-chat", @@ -837,6 +839,7 @@ test("cloud api AgentRun adapter reports persistent thread resume when a complet HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_ENVIRONMENT: "v02", @@ -976,6 +979,7 @@ test("cloud api AgentRun adapter exposes invalid tool-call attribution in result HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_ENVIRONMENT: "v02", @@ -1124,6 +1128,7 @@ test("cloud api AgentRun adapter maps minimax-m3 provider profile to AgentRun ba HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_ENVIRONMENT: "v02", @@ -1295,6 +1300,7 @@ test("cloud api AgentRun adapter scopes AgentRun sessions by backend profile", a HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek", HWLAB_ENVIRONMENT: "v02", @@ -1359,6 +1365,7 @@ test("cloud api AgentRun adapter rejects non-internal manager URLs by default", const server = createCloudApiServer({ env: { HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", AGENTRUN_MGR_URL: "http://74.48.78.17:8080" } }); @@ -1707,6 +1714,7 @@ test("cloud api trace uses AgentRun command result evidence when live trace stor HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, @@ -1823,6 +1831,7 @@ test("cloud api trace replays an earlier AgentRun command after same-run lastSeq HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, @@ -1911,6 +1920,7 @@ test("cloud api running trace refreshes AgentRun events without waiting for comm HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, @@ -2032,6 +2042,7 @@ test("cloud api repairs historical same-session AgentRun trace after lastTraceId HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, @@ -2190,6 +2201,7 @@ test("cloud api result polling repairs polluted completed AgentRun memory cache HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, @@ -2280,6 +2292,7 @@ test("cloud api result polling fails closed when AgentRun command registry misse HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", HWLAB_ENVIRONMENT: "v02", HWLAB_GITOPS_PROFILE: "v02" }, diff --git a/internal/cloud/server-code-agent-http.ts b/internal/cloud/server-code-agent-http.ts index 1cdf1fa3..4ddc0b29 100644 --- a/internal/cloud/server-code-agent-http.ts +++ b/internal/cloud/server-code-agent-http.ts @@ -128,7 +128,7 @@ export async function handleCodeAgentChatHttp(request, response, options) { } const payload = await runCodeAgentChat(nativeSessionChatParams, options); - await recordCodeAgentBillingUsage({ payload, params: nativeSessionChatParams, options }); + await finalizeCodeAgentBillingUsage({ payload, params: nativeSessionChatParams, options }); await recordCodeAgentSessionOwner({ payload, params: nativeSessionChatParams, options, status: payload.status === "completed" ? "active" : payload.status }); const responsePayload = annotateOwner(payload, nativeSessionChatParams); @@ -695,7 +695,7 @@ function submitCodeAgentChatTurn({ params, options, traceId }) { const traceStore = options.traceStore ?? defaultCodeAgentTraceStore; const results = options.codeAgentChatResults ?? createCodeAgentChatResultStore(); if (codeAgentAgentRunAdapterEnabled(options.env ?? process.env)) { - const initial = initialAgentRunChatResult({ params, options, traceId }); + const initial = withCodeAgentBillingReservation(initialAgentRunChatResult({ params, options, traceId }), params); results.set(traceId, annotateOwner(initial, params)); const run = async () => { let executionOptions = options; @@ -703,8 +703,7 @@ function submitCodeAgentChatTurn({ params, options, traceId }) { executionOptions = { ...options, ...(await codeAgentChatExecutionOptions(options, params)) }; const payload = await submitAgentRunChatTurn({ params, options: executionOptions, traceId, traceStore, results }); if (isCodeAgentResultCanceled(results.get(traceId))) return; - await recordCodeAgentBillingUsage({ payload, params, options: executionOptions }); - const owned = annotateOwner(payload, params); + const owned = annotateOwner(withCodeAgentBillingReservation(payload, params), params); await recordCodeAgentSessionOwner({ payload: owned, params, options: executionOptions, status: "running" }); results.set(traceId, owned); } catch (error) { @@ -729,7 +728,7 @@ function submitCodeAgentChatTurn({ params, options, traceId }) { updatedAt: new Date().toISOString() }, params); recordCodeAgentConversationFact(payload, executionOptions); - await recordCodeAgentBillingUsage({ payload, params, options: executionOptions }); + await finalizeCodeAgentBillingUsage({ payload, params, options: executionOptions }); results.set(traceId, payload); traceStore.append(traceId, { type: "result", @@ -772,7 +771,7 @@ function submitCodeAgentChatTurn({ params, options, traceId }) { try { const payload = await runCodeAgentChat(params, options); if (isCodeAgentResultCanceled(results.get(traceId))) return; - await recordCodeAgentBillingUsage({ payload, params, options }); + await finalizeCodeAgentBillingUsage({ payload, params, options }); await recordCodeAgentSessionOwner({ payload, params, options, status: payload.status === "completed" ? "active" : payload.status }); results.set(traceId, annotateOwner(payload, params)); traceStore.append(traceId, { @@ -802,7 +801,7 @@ function submitCodeAgentChatTurn({ params, options, traceId }) { }, updatedAt: new Date().toISOString() }; - await recordCodeAgentBillingUsage({ payload, params, options }); + await finalizeCodeAgentBillingUsage({ payload, params, options }); results.set(traceId, payload); traceStore.append(traceId, { type: "result", @@ -868,9 +867,10 @@ async function preflightCodeAgentBilling({ params = {}, options = {}, traceId, r } async function recordCodeAgentBillingUsage({ payload = {}, params = {}, options = {} } = {}) { - const reservation = params.userBillingReservation; + const reservation = params.userBillingReservation ?? payload.userBillingReservation; const reservationId = typeof reservation?.reservationId === "string" ? reservation.reservationId : ""; const client = options.userBillingClient; + if (payload.billing?.recorded === true || payload.billing?.released === true || payload.status === "running") return payload.billing ?? null; if (!reservationId || !client?.configured) return null; const traceId = safeTraceId(payload.traceId ?? params.traceId); const usedTokens = codeAgentUsedTokens(payload); @@ -905,6 +905,64 @@ async function recordCodeAgentBillingUsage({ payload = {}, params = {}, options return billing; } +async function finalizeCodeAgentBillingUsage({ payload = {}, params = {}, options = {} } = {}) { + if (payload.billing?.recorded === true || payload.billing?.released === true) return payload.billing; + if (payload.status === "running") return null; + if (payload.status === "completed") { + return recordCodeAgentBillingUsage({ payload, params, options }); + } + return releaseCodeAgentBillingReservation({ payload, params, options }); +} + +async function releaseCodeAgentBillingReservation({ payload = {}, params = {}, options = {} } = {}) { + const reservation = params.userBillingReservation ?? payload.userBillingReservation; + const reservationId = typeof reservation?.reservationId === "string" ? reservation.reservationId : ""; + const client = options.userBillingClient; + if (payload.billing?.recorded === true || payload.billing?.released === true || !reservationId || !client?.configured) return payload.billing ?? null; + const traceId = safeTraceId(payload.traceId ?? params.traceId); + const releaseReason = String(payload.status ?? "not_completed").trim() || "not_completed"; + const releaseBody = { + reservationId, + serviceId: "hwlab-code-agent", + idempotencyKey: `code-agent:${traceId}:release`, + reason: releaseReason, + metadata: codeAgentBillingMetadata(params, traceId, { + stage: "release", + status: payload.status ?? "unknown", + errorCode: payload.error?.code ?? payload.blocker?.code ?? null + }) + }; + const result = typeof client.billingRelease === "function" + ? await client.billingRelease(releaseBody) + : { ok: false, error: { code: "user_billing_release_not_supported" } }; + const billing = result.ok ? sanitizeBillingRelease(result.body, reservation) : { + released: false, + reservationId, + errorCode: result.error?.code ?? "user_billing_release_failed", + valuesRedacted: true + }; + payload.billing = billing; + const traceStore = options.traceStore ?? defaultCodeAgentTraceStore; + if (traceId) { + traceStore.append(traceId, { + type: "billing", + status: result.ok ? "released" : "degraded", + label: result.ok ? "billing:released" : "billing:release_failed", + reservationId, + reason: releaseReason, + errorCode: result.ok ? null : billing.errorCode, + valuesPrinted: false + }); + } + return billing; +} + +function withCodeAgentBillingReservation(payload = {}, params = {}) { + if (!payload || typeof payload !== "object" || payload.userBillingReservation) return payload; + const reservation = params.userBillingReservation; + return reservation ? { ...payload, userBillingReservation: reservation } : payload; +} + function codeAgentBillingEnabled(options = {}) { const env = options.env ?? process.env; if (String(env.HWLAB_USER_BILLING_CODE_AGENT_ENABLED ?? "").trim() === "0") return false; @@ -946,6 +1004,16 @@ function sanitizeBillingRecord(value = {}, reservation = {}) { }; } +function sanitizeBillingRelease(value = {}, reservation = {}) { + return { + released: true, + reservationId: textValue(value.reservationId) || textValue(reservation.reservationId) || null, + releasedCredits: Number.isFinite(Number(value.releasedCredits)) ? Number(value.releasedCredits) : null, + status: textValue(value.status) || "cancelled", + valuesRedacted: true + }; +} + function codeAgentUsedTokens(payload = {}) { const usage = payload.usage && typeof payload.usage === "object" ? payload.usage : null; const traceSummary = payload.traceSummary && typeof payload.traceSummary === "object" ? payload.traceSummary : null; @@ -1001,6 +1069,7 @@ export async function handleCodeAgentChatResultHttp(request, response, url, opti return; } if (synced.result && synced.result.status !== "running") { + await finalizeCodeAgentBillingUsage({ payload: synced.result, params: synced.result, options }); recordCodeAgentConversationFact(synced.result, options); await recordCodeAgentSessionOwner({ payload: synced.result, params: synced.result, options, status: codeAgentOwnerStatusForResult(synced.result) }); sendJson(response, 200, compactCodeAgentChatResultPayload(synced.result, options)); @@ -1761,6 +1830,7 @@ function codeAgentTraceResultEvidence(payload = {}, params = {}, traceId = null, const resolvedTraceId = safeTraceId(traceId); if (!resolvedTraceId) return null; const agentRun = agentRunSessionEvidence(payload).agentRun ?? null; + const userBillingReservation = codeAgentBillingReservationEvidence(payload.userBillingReservation ?? params.userBillingReservation); const resultSession = payload.session && typeof payload.session === "object" ? payload.session : null; const sessionReuse = payload.sessionReuse && typeof payload.sessionReuse === "object" ? payload.sessionReuse : null; const conversationId = safeConversationId(payload.conversationId ?? agentRun?.conversationId ?? resultSession?.conversationId ?? sessionReuse?.conversationId ?? params.conversationId) || null; @@ -1779,6 +1849,21 @@ function codeAgentTraceResultEvidence(payload = {}, params = {}, traceId = null, finalResponse, traceSummary, agentRun, + ...(userBillingReservation ? { userBillingReservation } : {}), + valuesRedacted: true, + secretMaterialStored: false + }; +} + +function codeAgentBillingReservationEvidence(value = null) { + if (!value || typeof value !== "object") return null; + const reservationId = textValue(value.reservationId); + if (!reservationId) return null; + return { + reservationId, + estimatedCredits: numberOrNull(value.estimatedCredits), + estimatedTokens: numberOrNull(value.estimatedTokens), + expiresAt: textValue(value.expiresAt) || null, valuesRedacted: true, secretMaterialStored: false }; @@ -2061,6 +2146,7 @@ export async function handleCodeAgentTraceHttp(request, response, url, options) agentRunResult = synced.result ?? agentRunResult; } if (isTraceCommandTerminalStatus(agentRunResult?.status)) { + await finalizeCodeAgentBillingUsage({ payload: agentRunResult, params: agentRunResult, options }); recordCodeAgentConversationFact(agentRunResult, options); await recordCodeAgentSessionOwner({ payload: agentRunResult, params: agentRunResult, options, status: codeAgentOwnerStatusForResult(agentRunResult), preserveLastTraceId: true }); } @@ -2298,11 +2384,12 @@ function compactCodeAgentChatResultPayload(payload, options = {}) { if (!payload || typeof payload !== "object") return payload; const limit = resultTraceEventLimit(options); const terminalEvidence = agentRunTerminalTraceEvidence(payload, payload.traceId); + const { userBillingReservation, ...publicPayload } = payload; return { - ...payload, + ...publicPayload, ...(terminalEvidence ? { terminalEvidence: terminalEvidencePayload(terminalEvidence) } : {}), - ...(payload.runnerTrace && typeof payload.runnerTrace === "object" - ? { runnerTrace: compactRunnerTraceForResult(payload.runnerTrace, limit) } + ...(publicPayload.runnerTrace && typeof publicPayload.runnerTrace === "object" + ? { runnerTrace: compactRunnerTraceForResult(publicPayload.runnerTrace, limit) } : {}) }; } diff --git a/internal/cloud/user-billing-client.ts b/internal/cloud/user-billing-client.ts index dffb8343..78146c89 100644 --- a/internal/cloud/user-billing-client.ts +++ b/internal/cloud/user-billing-client.ts @@ -90,6 +90,9 @@ export function createUserBillingClient({ env = process.env, fetchImpl = fetch } async billingRecord(body) { return post("/internal/billing/record", body); }, + async billingRelease(body) { + return post("/internal/billing/release", body); + }, async billingSummary(token, { limit = 20 } = {}) { const query = new URLSearchParams({ limit: String(limit) }); return requestJson(`/v1/billing/summary?${query.toString()}`, { method: "GET", bearerToken: token }); diff --git a/internal/cloud/user-billing-integration.test.ts b/internal/cloud/user-billing-integration.test.ts index 4fedff05..01a4be27 100644 --- a/internal/cloud/user-billing-integration.test.ts +++ b/internal/cloud/user-billing-integration.test.ts @@ -1,10 +1,13 @@ import assert from "node:assert/strict"; import { mkdtemp, rm } from "node:fs/promises"; +import { createServer as createHttpServer } from "node:http"; import os from "node:os"; import path from "node:path"; import { test } from "bun:test"; import { createCloudApiServer } from "./server.ts"; +import { createCodeAgentTraceStore } from "./code-agent-trace-store.ts"; +import { loadPersistedAgentRunResult } from "./code-agent-agentrun-adapter.ts"; import { codexStdioChatFixture, codexStdioReadyFixture } from "./server-test-helpers.ts"; test("cloud api accepts user-billing API keys and records Code Agent billing usage", async () => { @@ -201,6 +204,206 @@ test("cloud api accepts user-billing API keys and records Code Agent billing usa } }); +test("cloud api defers AgentRun Code Agent billing record until terminal result", async () => { + const calls = []; + const agentRunCalls = []; + let resultPolls = 0; + const agentRunServer = createHttpServer(async (request, response) => { + const url = new URL(request.url || "/", "http://127.0.0.1"); + const chunks = []; + for await (const chunk of request) chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); + const body = chunks.length ? JSON.parse(Buffer.concat(chunks).toString("utf8")) : null; + agentRunCalls.push({ method: request.method, path: url.pathname, search: url.search, body }); + const send = (data) => { + response.writeHead(200, { "content-type": "application/json" }); + response.end(`${JSON.stringify({ ok: true, data, traceId: "trc_fake_agentrun_billing" })}\n`); + }; + if (request.method === "POST" && url.pathname === "/api/v1/runs") { + assert.equal(body.backendProfile, "deepseek"); + return send({ id: "run_billing_deferred", status: "pending", backendProfile: "deepseek", sessionRef: body.sessionRef, resourceBundleRef: body.resourceBundleRef }); + } + if (request.method === "POST" && url.pathname === "/api/v1/runs/run_billing_deferred/commands") { + assert.equal(body.type, "turn"); + assert.equal(body.idempotencyKey, "trc_user_billing_agentrun_terminal"); + return send({ id: "cmd_billing_deferred", runId: "run_billing_deferred", state: "pending", type: "turn", seq: 1 }); + } + if (request.method === "GET" && url.pathname === "/api/v1/runs/run_billing_deferred/commands") { + return send({ items: [{ id: "cmd_billing_deferred", runId: "run_billing_deferred", state: "running", type: "turn", seq: 1, idempotencyKey: "trc_user_billing_agentrun_terminal", payload: { traceId: "trc_user_billing_agentrun_terminal", conversationId: "cnv_user_billing_agentrun", hwlabSessionId: "ses_user_billing_agentrun", providerProfile: "deepseek" } }] }); + } + if (request.method === "POST" && url.pathname === "/api/v1/runs/run_billing_deferred/runner-jobs") { + assert.equal(body.commandId, "cmd_billing_deferred"); + return send({ + action: "create-kubernetes-job", + runId: "run_billing_deferred", + commandId: "cmd_billing_deferred", + attemptId: "attempt_billing_deferred", + runnerId: "runner_billing_deferred", + namespace: "agentrun-v02", + jobName: "agentrun-v01-runner-billing-deferred" + }); + } + if (request.method === "GET" && url.pathname === "/api/v1/runs/run_billing_deferred/events") { + return send({ items: resultPolls > 0 ? [{ id: "evt_done", runId: "run_billing_deferred", seq: 1, type: "terminal_status", payload: { commandId: "cmd_billing_deferred", terminalStatus: "completed" }, createdAt: "2026-06-14T04:20:00.000Z" }] : [] }); + } + if (request.method === "GET" && url.pathname === "/api/v1/runs/run_billing_deferred/commands/cmd_billing_deferred/result") { + resultPolls += 1; + if (resultPolls === 1) { + return send({ runId: "run_billing_deferred", commandId: "cmd_billing_deferred", status: "running", runStatus: "claimed", commandState: "running", terminalStatus: null }); + } + return send({ + runId: "run_billing_deferred", + commandId: "cmd_billing_deferred", + attemptId: "attempt_billing_deferred", + runnerId: "runner_billing_deferred", + jobName: "agentrun-v01-runner-billing-deferred", + namespace: "agentrun-v02", + status: "completed", + runStatus: "completed", + commandState: "completed", + terminalStatus: "completed", + completed: true, + reply: "AgentRun billing completed.", + lastSeq: 1, + eventCount: 1, + sessionRef: { sessionId: "ses_agentrun_deepseek_billing", conversationId: "cnv_user_billing_agentrun", threadId: "thr_billing" } + }); + } + response.writeHead(404, { "content-type": "application/json" }); + response.end(`${JSON.stringify({ ok: false, failureKind: "unexpected", message: `${request.method} ${url.pathname}` })}\n`); + }); + await new Promise((resolve) => agentRunServer.listen(0, "127.0.0.1", resolve)); + + const userBillingClient = { + configured: true, + async introspect(token) { + calls.push({ op: "introspect", tokenPrefix: token.slice(0, 8) }); + assert.equal(token, "hwl_user_billing_agentrun_secret"); + return { ok: true, status: 200, body: { active: true, principal: { userId: "usr_user_billing_agentrun", email: "agentrun@hwlab.local", username: "billing-agentrun", role: "user", scopes: ["api"], authType: "api-key", keyId: "key_user_billing_agentrun" } } }; + }, + async billingPreflight(body) { + calls.push({ op: "preflight", body }); + assert.equal(body.apiKey, "hwl_user_billing_agentrun_secret"); + assert.equal(body.idempotencyKey, "code-agent:trc_user_billing_agentrun_terminal:preflight"); + return { ok: true, status: 200, body: { allowed: true, reservationId: "res_user_billing_agentrun", estimatedCredits: 1, expiresAt: "2026-06-14T05:00:00.000Z" } }; + }, + async billingRecord(body) { + calls.push({ op: "record", body }); + assert.equal(body.reservationId, "res_user_billing_agentrun"); + assert.equal(body.idempotencyKey, "code-agent:trc_user_billing_agentrun_terminal:record"); + assert.equal(body.metadata.status, "completed"); + return { ok: true, status: 200, body: { recordId: "use_user_billing_agentrun", credits: 1, balance: 9 } }; + }, + async billingRelease() { + throw new Error("completed AgentRun billing should record usage instead of releasing reservation"); + } + }; + const { port: agentRunPort } = agentRunServer.address(); + const server = createCloudApiServer({ + traceStore: createCodeAgentTraceStore(), + env: { + HWLAB_ACCESS_CONTROL_REQUIRED: "1", + HWLAB_USER_BILLING_CODE_AGENT_ENABLED: "1", + HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01", + AGENTRUN_MGR_URL: `http://127.0.0.1:${agentRunPort}`, + AGENTRUN_API_KEY: "test-agentrun-key", + HWLAB_CODE_AGENT_AGENTRUN_ALLOW_NON_K3S_URL: "1", + HWLAB_CODE_AGENT_AGENTRUN_SOURCE_COMMIT: "0123456789abcdef0123456789abcdef01234567", + HWLAB_CODE_AGENT_AGENTRUN_PROVIDER_ID: "D601", + HWLAB_CODE_AGENT_AGENTRUN_SESSION_STORAGE: "metadata-only", + HWLAB_CODE_AGENT_DEFAULT_PROVIDER_PROFILE: "deepseek" + }, + userBillingClient + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + + try { + const { port } = server.address(); + const authHeader = { authorization: "Bearer hwl_user_billing_agentrun_secret" }; + const sessionResponse = await fetch(`http://127.0.0.1:${port}/v1/agent/sessions`, { + method: "POST", + headers: { "content-type": "application/json", ...authHeader }, + body: JSON.stringify({ conversationId: "cnv_user_billing_agentrun", sessionId: "ses_user_billing_agentrun", providerProfile: "deepseek" }) + }); + assert.equal(sessionResponse.status, 201); + + const submit = await fetch(`http://127.0.0.1:${port}/v1/agent/chat`, { + method: "POST", + headers: { "content-type": "application/json", authorization: "Bearer hwl_user_billing_agentrun_secret", "x-trace-id": "trc_user_billing_agentrun_terminal", prefer: "respond-async" }, + body: JSON.stringify({ conversationId: "cnv_user_billing_agentrun", sessionId: "ses_user_billing_agentrun", shortConnection: true, message: "billing AgentRun terminal smoke" }) + }); + assert.equal(submit.status, 202); + await waitForUserBillingCondition(() => agentRunCalls.some((call) => call.path === "/api/v1/runs/run_billing_deferred/runner-jobs")); + assert.deepEqual(calls.map((call) => call.op), ["introspect", "introspect", "preflight"]); + + const running = await fetch(`http://127.0.0.1:${port}/v1/agent/chat/result/trc_user_billing_agentrun_terminal`, { headers: authHeader }); + assert.equal(running.status === 202 || running.status === 200, true); + assert.equal(calls.some((call) => call.op === "record"), false); + + const payload = running.status === 200 ? await running.json() : await pollUserBillingAgentResult(port, "trc_user_billing_agentrun_terminal", authHeader); + assert.equal(payload.status, "completed"); + assert.equal(payload.billing.recorded, true); + assert.equal(payload.billing.reservationId, "res_user_billing_agentrun"); + assert.equal(JSON.stringify(payload).includes("userBillingReservation"), false); + assert.equal(calls.filter((call) => call.op === "preflight").length, 1); + assert.equal(calls.filter((call) => call.op === "record").length, 1); + assert.equal(calls.some((call) => call.op === "release"), false); + } finally { + await new Promise((resolve, reject) => server.close((error) => (error ? reject(error) : resolve()))); + await new Promise((resolve, reject) => agentRunServer.close((error) => (error ? reject(error) : resolve()))); + } +}); + +test("AgentRun persisted trace evidence restores billing reservation without exposing it on agentRun", async () => { + const restored = await loadPersistedAgentRunResult("trc_user_billing_agentrun_restored", { + env: { HWLAB_CODE_AGENT_ADAPTER: "agentrun-v01" }, + accessController: { + async getAgentSessionByTraceId(traceId) { + assert.equal(traceId, "trc_user_billing_agentrun_restored"); + return { + id: "ses_user_billing_agentrun_restored", + ownerUserId: "usr_user_billing_agentrun", + ownerRole: "user", + conversationId: "cnv_user_billing_agentrun", + threadId: "thr_billing_restored", + status: "running", + startedAt: "2026-06-14T04:00:00.000Z", + updatedAt: "2026-06-14T04:10:00.000Z", + session: { + traceResults: { + trc_user_billing_agentrun_restored: { + traceId: "trc_user_billing_agentrun_restored", + status: "running", + userBillingReservation: { + reservationId: "res_user_billing_restored", + estimatedCredits: 1, + expiresAt: "2026-06-14T05:00:00.000Z", + valuesRedacted: true, + secretMaterialStored: false + }, + agentRun: { + adapter: "agentrun-v01", + managerUrl: "http://agentrun-mgr.agentrun-v02.svc.cluster.local:8080", + backendProfile: "deepseek", + providerId: "D601", + runId: "run_billing_restored", + commandId: "cmd_billing_restored", + status: "running", + commandState: "running", + traceId: "trc_user_billing_agentrun_restored", + valuesPrinted: false + } + } + } + } + }; + } + } + }); + assert.equal(restored.userBillingReservation.reservationId, "res_user_billing_restored"); + assert.equal(restored.agentRun.runId, "run_billing_restored"); + assert.equal(Object.hasOwn(restored.agentRun, "userBillingReservation"), false); +}); + async function pollUserBillingAgentResult(port, traceId, headers) { let last = null; for (let attempt = 0; attempt < 20; attempt += 1) { @@ -215,6 +418,15 @@ async function pollUserBillingAgentResult(port, traceId, headers) { throw new Error(`Code Agent result did not complete: ${JSON.stringify(last)}`); } +async function waitForUserBillingCondition(predicate, timeoutMs = 500) { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) return; + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error("condition was not met before timeout"); +} + test("cloud api proxies admin billing summary for web admins without exposing user-billing secrets", async () => { const calls = []; const userBillingClient = { diff --git a/internal/userbilling/service.go b/internal/userbilling/service.go index 01ec3997..0c09de2e 100644 --- a/internal/userbilling/service.go +++ b/internal/userbilling/service.go @@ -284,6 +284,7 @@ func (s *Server) routes() { s.route(http.MethodPost, "/internal/auth/introspect", s.handleIntrospect) s.route(http.MethodPost, "/internal/billing/preflight", s.handleBillingPreflight) s.route(http.MethodPost, "/internal/billing/record", s.handleBillingRecord) + s.route(http.MethodPost, "/internal/billing/release", s.handleBillingRelease) s.route(http.MethodGet, "/internal/admin/billing/summary", s.handleAdminBillingSummary) s.route(http.MethodPost, "/internal/admin/credits/adjust", s.handleAdminCreditAdjust) s.mux.HandleFunc("/internal/admin/users", s.handleAdminUsers) @@ -936,6 +937,30 @@ func (s *Server) handleBillingRecord(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, result) } +func (s *Server) handleBillingRelease(w http.ResponseWriter, r *http.Request) { + if !s.internalAllowed(w, r) || !s.databaseAvailable(w) { + return + } + var req struct { + ReservationID string `json:"reservationId"` + ServiceID string `json:"serviceId"` + Reason string `json:"reason"` + IdempotencyKey string `json:"idempotencyKey"` + Metadata map[string]any `json:"metadata"` + } + if !decodeJSON(w, r, &req) { + return + } + ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) + defer cancel() + result, err := s.releaseReservation(ctx, req.ReservationID, req.ServiceID, req.Reason, req.IdempotencyKey, req.Metadata) + if err != nil { + writeAPIError(w, http.StatusBadRequest, "billing_release_failed", err.Error()) + return + } + writeJSON(w, http.StatusOK, result) +} + func (s *Server) handleAdminUsers(w http.ResponseWriter, r *http.Request) { if !s.internalAllowed(w, r) || !s.databaseAvailable(w) { return @@ -1404,6 +1429,62 @@ func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID return map[string]any{"recordId": recordID, "userId": userID, "serviceId": serviceID, "credits": credits, "balance": newBalance}, nil } +func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceID, reason, idempotencyKey string, metadata map[string]any) (map[string]any, error) { + reservationID = strings.TrimSpace(reservationID) + if reservationID == "" { + return nil, errors.New("reservationId is required") + } + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{}) + if err != nil { + return nil, err + } + defer tx.Rollback() + var userID, actualServiceID, status string + var reserved int64 + err = tx.QueryRowContext(ctx, `SELECT user_id, service_id, estimated_credits, status FROM hwlab_billing_reservations WHERE id = $1 FOR UPDATE`, reservationID).Scan(&userID, &actualServiceID, &reserved, &status) + if err != nil { + return nil, err + } + serviceID = strings.TrimSpace(serviceID) + if serviceID != "" && serviceID != actualServiceID { + return nil, fmt.Errorf("reservation service is %s", actualServiceID) + } + if status == "cancelled" || status == "expired" { + return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "releasedCredits": int64(0), "status": status, "idempotent": true}, tx.Commit() + } + if status != "reserved" { + return nil, fmt.Errorf("reservation status is %s", status) + } + var accountReserved int64 + if err := tx.QueryRowContext(ctx, `SELECT reserved_credits FROM hwlab_credit_accounts WHERE user_id = $1 FOR UPDATE`, userID).Scan(&accountReserved); err != nil { + return nil, err + } + newReserved := accountReserved - reserved + if newReserved < 0 { + newReserved = 0 + } + _, err = tx.ExecContext(ctx, `UPDATE hwlab_credit_accounts SET reserved_credits = $2, updated_at = now() WHERE user_id = $1`, userID, newReserved) + if err != nil { + return nil, err + } + releaseMetadata := map[string]any{ + "releaseReason": strings.TrimSpace(reason), + "releaseIdempotency": strings.TrimSpace(idempotencyKey), + "releaseValuesRedacted": true, + } + for key, value := range metadata { + releaseMetadata[key] = value + } + _, err = tx.ExecContext(ctx, `UPDATE hwlab_billing_reservations SET status = 'cancelled', metadata = COALESCE(metadata, '{}'::jsonb) || $2::jsonb, updated_at = now() WHERE id = $1`, reservationID, jsonObject(releaseMetadata)) + if err != nil { + return nil, err + } + if err := tx.Commit(); err != nil { + return nil, err + } + return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "releasedCredits": reserved, "status": "cancelled"}, nil +} + func (s *Server) adjustCredits(ctx context.Context, userID string, delta int64, kind, reason, idempotencyKey string, metadata map[string]any) (int64, error) { tx, err := s.db.BeginTx(ctx, &sql.TxOptions{}) if err != nil {