diff --git a/billing-token-state.mjs b/billing-token-state.mjs index 5408fef..e5d6610 100644 --- a/billing-token-state.mjs +++ b/billing-token-state.mjs @@ -1,4 +1,5 @@ import { normalizeTokenState } from './billing.mjs'; +import { fetchGooseSessionAccumulatedCostUsd } from './goose-session-cost.mjs'; export function loadCostEstimateConfig(env = process.env) { const useBackendCost = env.H5_USE_BACKEND_COST === '1'; @@ -56,22 +57,44 @@ export function enrichTokenStateForBilling( return state; } +async function resolveSessionCostPayload(sessionId, fetchSession, fetchSessionCostFromPg, env) { + let sessionCost = null; + if (typeof fetchSession === 'function' && sessionId) { + try { + sessionCost = await fetchSession(sessionId); + } catch { + sessionCost = null; + } + } + if (pickSessionAccumulatedCost(sessionCost) != null) { + return sessionCost; + } + + const readPgCost = + typeof fetchSessionCostFromPg === 'function' + ? fetchSessionCostFromPg + : (sid) => fetchGooseSessionAccumulatedCostUsd(sid, env); + const pgUsd = env.H5_USE_BACKEND_COST === '1' ? await readPgCost(sessionId) : null; + if (pgUsd != null) { + return { + ...(sessionCost && typeof sessionCost === 'object' ? sessionCost : {}), + accumulated_cost: pgUsd, + }; + } + return sessionCost; +} + export async function resolveBillingTokenState( tokenStateRaw, - { sessionId = null, fetchSession = null } = {}, + { sessionId = null, fetchSession = null, fetchSessionCostFromPg = null } = {}, env = process.env, ) { const state = normalizeTokenState(tokenStateRaw); - - if (typeof fetchSession === 'function' && sessionId) { - try { - const session = await fetchSession(sessionId); - const enriched = enrichTokenStateForBilling(state, { sessionCost: session }, env); - if (enriched.accumulatedCost != null) return enriched; - } catch { - // Best-effort: fall through to inline cost / token estimate. - } - } - - return enrichTokenStateForBilling(state, {}, env); + const sessionCost = await resolveSessionCostPayload( + sessionId, + fetchSession, + fetchSessionCostFromPg, + env, + ); + return enrichTokenStateForBilling(state, { sessionCost }, env); } diff --git a/billing-token-state.test.mjs b/billing-token-state.test.mjs index 7c437ed..4ef7035 100644 --- a/billing-token-state.test.mjs +++ b/billing-token-state.test.mjs @@ -103,6 +103,23 @@ test('resolveBillingTokenState replaces inflated Finish cost with session cost', assert.equal(resolved.accumulatedCost, 0.006738208); }); +test('resolveBillingTokenState uses PG cost when goosed session API omits accumulated_cost', async () => { + const resolved = await resolveBillingTokenState( + { + accumulatedInputTokens: 280030, + accumulatedOutputTokens: 5797, + accumulatedCost: 0.082, + }, + { + sessionId: '20260804_19', + fetchSession: async () => ({ id: '20260804_19', accumulated_cost: null }), + fetchSessionCostFromPg: async () => 0.006738208, + }, + { H5_USE_BACKEND_COST: '1' }, + ); + assert.equal(resolved.accumulatedCost, 0.006738208); +}); + test('enriched cost drives 1.2x billing instead of flat fallback', () => { const previous = { lastInputTokens: 224853, lastOutputTokens: 6685, lastAccumulatedCost: null }; const tokenState = enrichTokenStateForBilling( diff --git a/goose-session-cost.mjs b/goose-session-cost.mjs new file mode 100644 index 0000000..06261e1 --- /dev/null +++ b/goose-session-cost.mjs @@ -0,0 +1,71 @@ +let sharedPgClientPromise = null; + +export function resolveGooseSessionPgUrl(env = process.env) { + const explicit = String(env.GOOSE_SESSION_DB_URL ?? '').trim(); + if (explicit) return explicit; + const host = env.GOOSE_SESSION_PG_HOST ?? '127.0.0.1'; + const port = env.GOOSE_SESSION_PG_PORT ?? '5432'; + const database = env.GOOSE_SESSION_PG_DATABASE ?? 'memind_sessions'; + const user = env.GOOSE_SESSION_PG_USER ?? 'john'; + const password = env.GOOSE_SESSION_PG_PASSWORD ?? ''; + if (password) { + return `postgresql://${encodeURIComponent(user)}:${encodeURIComponent(password)}@${host}:${port}/${database}`; + } + return `postgresql://${user}@${host}:${port}/${database}`; +} + +export function isGooseSessionPgConfigured(env = process.env) { + if (String(env.GOOSE_SESSION_DB_URL ?? '').trim()) return true; + return String(env.GOOSE_SESSION_PG_DISABLE ?? '') !== '1'; +} + +async function getSharedPgClient(env = process.env) { + if (!isGooseSessionPgConfigured(env)) return null; + if (!sharedPgClientPromise) { + sharedPgClientPromise = import('pg') + .then(async ({ default: pg }) => { + const client = new pg.Client({ connectionString: resolveGooseSessionPgUrl(env) }); + await client.connect(); + return client; + }) + .catch((err) => { + sharedPgClientPromise = null; + throw err; + }); + } + return sharedPgClientPromise; +} + +export async function fetchGooseSessionAccumulatedCostUsd(sessionId, env = process.env) { + const normalizedSessionId = String(sessionId ?? '').trim(); + if (!normalizedSessionId || !isGooseSessionPgConfigured(env)) return null; + try { + const client = await getSharedPgClient(env); + const result = await client.query( + `SELECT accumulated_cost FROM sessions WHERE id = $1 LIMIT 1`, + [normalizedSessionId], + ); + const raw = result.rows[0]?.accumulated_cost; + if (raw == null) return null; + const value = Number(raw); + return Number.isFinite(value) && value >= 0 ? value : null; + } catch { + return null; + } +} + +export function createGooseSessionCostReader(env = process.env) { + return (sessionId) => fetchGooseSessionAccumulatedCostUsd(sessionId, env); +} + +export async function closeGooseSessionPgClient() { + if (!sharedPgClientPromise) return; + try { + const client = await sharedPgClientPromise; + await client.end(); + } catch { + // ignore shutdown errors in tests/scripts + } finally { + sharedPgClientPromise = null; + } +} diff --git a/goose-session-cost.test.mjs b/goose-session-cost.test.mjs new file mode 100644 index 0000000..be55cbc --- /dev/null +++ b/goose-session-cost.test.mjs @@ -0,0 +1,30 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { + isGooseSessionPgConfigured, + resolveGooseSessionPgUrl, +} from './goose-session-cost.mjs'; + +test('resolveGooseSessionPgUrl prefers GOOSE_SESSION_DB_URL', () => { + assert.equal( + resolveGooseSessionPgUrl({ GOOSE_SESSION_DB_URL: 'postgresql://u:p@host:5432/db' }), + 'postgresql://u:p@host:5432/db', + ); +}); + +test('resolveGooseSessionPgUrl builds local default DSN', () => { + assert.equal( + resolveGooseSessionPgUrl({ + GOOSE_SESSION_PG_HOST: '127.0.0.1', + GOOSE_SESSION_PG_PORT: '5432', + GOOSE_SESSION_PG_DATABASE: 'memind_sessions', + GOOSE_SESSION_PG_USER: 'john', + }), + 'postgresql://john@127.0.0.1:5432/memind_sessions', + ); +}); + +test('isGooseSessionPgConfigured can be disabled explicitly', () => { + assert.equal(isGooseSessionPgConfigured({ GOOSE_SESSION_PG_DISABLE: '1' }), false); + assert.equal(isGooseSessionPgConfigured({ GOOSE_SESSION_DB_URL: 'postgresql://x' }), true); +}); diff --git a/scripts/compensate-billing-token-estimate-overcharge.mjs b/scripts/compensate-billing-token-estimate-overcharge.mjs index 1069a73..4a1b2ef 100644 --- a/scripts/compensate-billing-token-estimate-overcharge.mjs +++ b/scripts/compensate-billing-token-estimate-overcharge.mjs @@ -13,6 +13,7 @@ import path from 'node:path'; import { fileURLToPath } from 'node:url'; import mysql from 'mysql2/promise'; import pg from 'pg'; +import { resolveGooseSessionPgUrl } from '../goose-session-cost.mjs'; const root = path.join(path.dirname(fileURLToPath(import.meta.url)), '..'); @@ -33,9 +34,14 @@ loadEnvFile(path.join(root, '.env')); const apply = process.argv.includes('--apply'); const sinceArg = process.argv.find((a) => a.startsWith('--since=')); -const sinceDate = sinceArg ? sinceArg.slice('--since='.length) : '2026-08-04'; -const startMs = new Date(`${sinceDate}T00:00:00+08:00`).getTime(); -const DEDupe_NOTE_PREFIX = `补偿：Token估价超扣(${sinceDate}起)`; +const sinceRaw = sinceArg ? sinceArg.slice('--since='.length) : '2026-08-04'; +const startMs = sinceRaw.includes('T') + ? new Date(sinceRaw).getTime() + : new Date(`${sinceRaw}T00:00:00+08:00`).getTime(); +const sinceLabel = sinceRaw.includes('T') + ? sinceRaw.replace('T', ' ').replace('+08:00', ' CST') + : `${sinceRaw} 00:00 CST`; +const DEDupe_NOTE_PREFIX = `补偿：Token估价超扣(${sinceLabel}起)`; function loadBillingConfig() { const marginMultiplier = Number(process.env.H5_MARGIN_MULTIPLIER ?? 1); @@ -56,12 +62,7 @@ function correctTotalCents(gooseCostUsd, config) { } function resolvePgUrl() { - if (process.env.GOOSE_SESSION_DB_URL) return process.env.GOOSE_SESSION_DB_URL; - const host = process.env.GOOSE_SESSION_PG_HOST ?? '127.0.0.1'; - const port = process.env.GOOSE_SESSION_PG_PORT ?? '5432'; - const db = process.env.GOOSE_SESSION_PG_DATABASE ?? 'memind_sessions'; - const user = process.env.GOOSE_SESSION_PG_USER ?? 'john'; - return `postgresql://${user}@${host}:${port}/${db}`; + return resolveGooseSessionPgUrl(process.env); } async function main() { @@ -102,7 +103,7 @@ async function main() { const sessionIds = [...bySession.keys()]; if (sessionIds.length === 0) { - console.log('No usage records since', sinceDate); + console.log('No usage records since', sinceLabel); return; } @@ -150,7 +151,7 @@ async function main() { const totalRefund = users.reduce((sum, user) => sum + user.refundCents, 0); console.log('Billing config:', config); - console.log('Since:', sinceDate, `(${startMs})`); + console.log('Since:', sinceLabel, `(${startMs})`); console.log('Affected users:', users.length); console.log('Total refund:', `¥${(totalRefund / 100).toFixed(2)}`, `(${totalRefund} cents)`); console.log('');