Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 134 additions & 0 deletions apps/desktop/src/main/__tests__/turnCostBroadcaster.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
* per-turn 费用挂载(MessageActionBar"本轮消耗")的 main 侧业务体:
* - recordTurnCostOnMessage:patch 成功才广播;patch false(行不存在)不广播;
* costUsd 非法 / 极小直接跳过(绝不写 $0);patch 抛错只吞不传播。
* - recordTurnUsageOnMessage:算不出报价的轮次只落 turnUsageDetails,不碰任何
* 金额字段与 scheduler 账本(UI 据此退回显示本轮 token)。
* - codexUsageToTokens:done.data.usage → computeGatewayTurnCost 入参映射
* (reasoning 算 output,与 daily_model_usage 口径一致)。
*
Expand Down Expand Up @@ -37,6 +39,7 @@ vi.mock('../messagePersistBroadcaster.js', () => ({
import {
recordTurnCostOnMessage,
recordSchedulerTurnCost,
recordTurnUsageOnMessage,
codexUsageToTokens,
type TurnCostDeps,
type MessageTurnCostPayload,
Expand Down Expand Up @@ -331,6 +334,137 @@ describe('recordTurnCostOnMessage', () => {
});
});

describe('recordTurnUsageOnMessage', () => {
const USAGE_ARGS = { sessionId: 's1', clientId: 'm1', turnUsageDetails: DETAILS };

it('只落 token 明细,不写任何金额字段', async () => {
const { deps, broadcasts, patchCalls, runCostCalls } = makeDeps(true);
await expect(recordTurnUsageOnMessage(USAGE_ARGS, deps)).resolves.toBe(true);
// 账本口径:没有钱就不碰钱。patch 里只能有 turnUsageDetails 一个键。
expect(patchCalls).toEqual([
{
sessionId: 's1',
clientId: 'm1',
patch: { turnUsageDetails: DETAILS },
},
]);
// scheduler 费用账本不参与(它只接受真实计费)。
expect(runCostCalls).toHaveLength(0);
expect(broadcasts).toEqual([
{ sessionId: 's1', clientId: 'm1', turnUsageDetails: DETAILS },
]);
});

it('广播 payload 不带金额字段 —— 消费方据此走 token 回退', async () => {
const { deps, broadcasts } = makeDeps(true);
await recordTurnUsageOnMessage(USAGE_ARGS, deps);
const [payload] = broadcasts;
expect(payload.turnMoney).toBeUndefined();
expect(payload.turnCostUsd).toBeUndefined();
expect(payload.userTurnMoney).toBeUndefined();
expect(payload.userTurnCostUsd).toBeUndefined();
});

it('明细缺省(整轮 0 token)→ 不落库不广播,绝不写空对象', async () => {
const { deps, broadcasts, patchCalls } = makeDeps(true);
await expect(
recordTurnUsageOnMessage({ sessionId: 's1', clientId: 'm1' }, deps),
).resolves.toBe(false);
expect(patchCalls).toHaveLength(0);
expect(broadcasts).toHaveLength(0);
});

it('sessionId / clientId 缺失 → 直接跳过', async () => {
const { deps, patchCalls } = makeDeps(true);
await expect(
recordTurnUsageOnMessage({ sessionId: '', clientId: 'm1', turnUsageDetails: DETAILS }, deps),
).resolves.toBe(false);
await expect(
recordTurnUsageOnMessage({ sessionId: 's1', clientId: '', turnUsageDetails: DETAILS }, deps),
).resolves.toBe(false);
expect(patchCalls).toHaveLength(0);
});

it('patch 返回 null(行已被 rewind 删)→ 不广播', async () => {
const { deps, broadcasts } = makeDeps(false);
await expect(recordTurnUsageOnMessage(USAGE_ARGS, deps)).resolves.toBe(false);
expect(broadcasts).toHaveLength(0);
});

it('patch 抛错 → 只吞不传播(调用方 fire-and-forget)', async () => {
const { deps, broadcasts } = makeDeps(new Error('db locked'));
await expect(recordTurnUsageOnMessage(USAGE_ARGS, deps)).resolves.toBe(false);
expect(broadcasts).toHaveLength(0);
});

// 一次用户请求可能含多个 SDK segment(自动续跑):前面的有真实费用、最后一个缺报价
// 走 usage-only 时,若只写明细,收尾消息会退回显示 token,把这一轮已经花掉的钱藏起来。
it('本轮此前已有费用 → 收尾消息带上用户轮累计,而不是退回 token', async () => {
const { deps, broadcasts, patchCalls, runCostCalls } = makeDeps(true, {
money: usdMoney(0.31),
costUsd: 0.31,
hasEstimatedValue: false,
});
await expect(recordTurnUsageOnMessage(USAGE_ARGS, deps)).resolves.toBe(true);

expect(patchCalls[0].patch).toEqual({
turnUsageDetails: DETAILS,
userTurnCost: usdMoney(0.31),
userTurnCostUsd: 0.31,
userTurnCostIsEstimate: false,
});
// 当前这个无价 segment 依然不记账:没有 turnCost / turnCostUsd,也不碰 scheduler 账本。
expect(patchCalls[0].patch).not.toHaveProperty('turnCost');
expect(patchCalls[0].patch).not.toHaveProperty('turnCostUsd');
expect(runCostCalls).toHaveLength(0);

expect(broadcasts[0]).toMatchObject({
userTurnMoney: usdMoney(0.31),
userTurnCostUsd: 0.31,
userTurnCostIsEstimate: false,
turnUsageDetails: DETAILS,
});
expect(broadcasts[0].turnMoney).toBeUndefined();
});

it('本轮此前的费用是订阅价值估算 → 估算标记一并带上', async () => {
const { deps, patchCalls, broadcasts } = makeDeps(true, {
money: usdMoney(0.2, 'value-estimate'),
costUsd: 0.2,
hasEstimatedValue: true,
});
await recordTurnUsageOnMessage(USAGE_ARGS, deps);
expect(patchCalls[0].patch).toMatchObject({ userTurnCostIsEstimate: true });
expect(broadcasts[0]).toMatchObject({ userTurnCostIsEstimate: true });
});

it('本轮此前没有费用 → 只写明细(不写 0 值累计)', async () => {
const { deps, patchCalls, broadcasts } = makeDeps(true);
await recordTurnUsageOnMessage(USAGE_ARGS, deps);
expect(patchCalls[0].patch).toEqual({ turnUsageDetails: DETAILS });
expect(broadcasts[0].userTurnMoney).toBeUndefined();
});

it('CNY 累计不写 userTurnCostUsd 字段', async () => {
const { deps, patchCalls, broadcasts } = makeDeps(true, {
money: cnyMoney(2.5),
costUsd: 0,
hasEstimatedValue: false,
});
await recordTurnUsageOnMessage(USAGE_ARGS, deps);
expect(patchCalls[0].patch).not.toHaveProperty('userTurnCostUsd');
expect(broadcasts[0].userTurnCostUsd).toBeUndefined();
expect(broadcasts[0].userTurnMoney).toEqual(cnyMoney(2.5));
});

it('读取往轮累计失败 → 不落库不广播(不猜金额)', async () => {
const { deps, patchCalls, broadcasts } = makeDeps(true, new Error('db locked'));
await expect(recordTurnUsageOnMessage(USAGE_ARGS, deps)).resolves.toBe(false);
expect(patchCalls).toHaveLength(0);
expect(broadcasts).toHaveLength(0);
});
});

describe('recordSchedulerTurnCost', () => {
const schedulerOrigin = {
kind: 'scheduler',
Expand Down
132 changes: 87 additions & 45 deletions apps/desktop/src/main/maker-ipc/register.ts
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,7 @@ import {
codexUsageToTokens,
recordSchedulerTurnCost,
recordTurnCostOnMessage,
recordTurnUsageOnMessage,
} from '../turnCostBroadcaster.js';
import { recordModelMismatchOnMessage } from '../modelMismatchBroadcaster.js';
import { detectClaudeModelMismatch } from '../../shared/modelMismatch.js';
Expand Down Expand Up @@ -3297,8 +3298,8 @@ export function wireSessionToIpc(session: ReturnType<Maker['getSession']>): void
estimatedValues.length > 0
? addRegionalMoney(estimatedValues)
: null;
const turnUsageDetails = buildClaudeTurnUsageDetails(doneData?.usage, deltas, 'unknown', perModel);
if (turnEstimatedValue && turnEstimatedValue.amount > 0) {
const turnUsageDetails = buildClaudeTurnUsageDetails(doneData?.usage, deltas, 'unknown', perModel);
const changedScheduleId = await recordSchedulerTurnCost({
sessionId: session.id,
clientId: turnAssistantPersistId,
Expand All @@ -3307,53 +3308,80 @@ export function wireSessionToIpc(session: ReturnType<Maker['getSession']>): void
turnOrigin: event.turnOrigin,
});
if (changedScheduleId) broadcastSchedulerChanged(changedScheduleId);
} else {
// 真实计费与订阅估值都拿不到(典型:网关目录整体不下发价格、模型不在价表)
// —— 钱没有,但 token 明细是算好的,落下来让 UI 退回显示本轮 token。
await recordTurnUsageOnMessage({
sessionId: session.id,
clientId: turnAssistantPersistId,
turnUsageDetails,
});
}
}
})();
} else if (typeof cumulative === 'number' && cumulative >= 0) {
// 窄兜底: 罕见地 done 只带 total_cost_usd、没 modelUsage —— 拆不了 daily_model_usage,
// 但至少用累计差把总额 / session / message 记上, 别漏整轮 (review #4)。
const rawDelta = Math.max(0, cumulative - prevReportedCost);
if (rawDelta > 0) {
void (async () => {
let resolvedModel = 'unknown';
try {
const model = await modelPromise;
resolvedModel = model;
} catch { /* non-fatal: 保留 SDK 原始 cost */ }
// 订阅直连轮(chatgpt/ / xai/)走窄兜底时: 真实计费恒 0, 不写 daily_spend /
// sessions.total_cost_usd(与主路径 resolveTurnCost 的 subscription gate 同口径,
// 避免把订阅 SDK 自报 cost 误记进计费)。
if (isSubscriptionDirectModel(resolvedModel)) return;
const providerId = getSessionProvider(session.id);
const observedRoute =
providerId == null ? readClaudeSessionRoute(session.id) : null;
const route: BillingRoute = session.remoteHostId
? 'unknown'
: providerId === 'anthropic' || observedRoute === 'subscription'
? 'subscription'
: providerId === 'xd' || observedRoute === 'gateway'
? 'xd-gateway'
: providerId
? 'provider-api'
: 'unknown';
if (route === 'subscription' || route === 'xd-gateway') return;
const ledgerCurrency =
(await getGatewayAccountCurrency()) ?? currentLedgerCurrency();
const money = usdToLedgerCurrency(rawDelta, ledgerCurrency);
const turnUsageDetails = buildClaudeTurnUsageDetails(doneData?.usage, undefined, resolvedModel);
recordTurnSpend(money);
recordSessionTurnSpend(session.id, money);
const changedScheduleId = await recordSchedulerTurnCost({
void (async () => {
let resolvedModel = 'unknown';
try {
const model = await modelPromise;
resolvedModel = model;
} catch { /* non-fatal: 保留 SDK 原始 cost */ }
const turnUsageDetails = buildClaudeTurnUsageDetails(doneData?.usage, undefined, resolvedModel);
// 本分支有三个"记不了钱"的出口(本轮 cost 未增长 / 订阅直连 / 订阅与网关路由),
// 账本口径一个字不改,但都把本轮 token 明细落下来 —— 钱算不出来不代表用量
// 算不出来,UI 那一格据此退回显示 token 而不是空着。
const recordUsageOnly = async () => {
if (!turnAssistantPersistId) return;
await recordTurnUsageOnMessage({
sessionId: session.id,
clientId: turnAssistantPersistId,
money,
turnUsageDetails,
turnOrigin: event.turnOrigin,
});
if (changedScheduleId) broadcastSchedulerChanged(changedScheduleId);
})();
}
};
if (rawDelta <= 0) {
await recordUsageOnly();
return;
}
// 订阅直连轮(chatgpt/ / xai/)走窄兜底时: 真实计费恒 0, 不写 daily_spend /
// sessions.total_cost_usd(与主路径 resolveTurnCost 的 subscription gate 同口径,
// 避免把订阅 SDK 自报 cost 误记进计费)。
if (isSubscriptionDirectModel(resolvedModel)) {
await recordUsageOnly();
return;
}
const providerId = getSessionProvider(session.id);
const observedRoute =
providerId == null ? readClaudeSessionRoute(session.id) : null;
const route: BillingRoute = session.remoteHostId
? 'unknown'
: providerId === 'anthropic' || observedRoute === 'subscription'
? 'subscription'
: providerId === 'xd' || observedRoute === 'gateway'
? 'xd-gateway'
: providerId
? 'provider-api'
: 'unknown';
if (route === 'subscription' || route === 'xd-gateway') {
await recordUsageOnly();
return;
}
const ledgerCurrency =
(await getGatewayAccountCurrency()) ?? currentLedgerCurrency();
const money = usdToLedgerCurrency(rawDelta, ledgerCurrency);
recordTurnSpend(money);
recordSessionTurnSpend(session.id, money);
const changedScheduleId = await recordSchedulerTurnCost({
sessionId: session.id,
clientId: turnAssistantPersistId,
money,
turnUsageDetails,
turnOrigin: event.turnOrigin,
});
if (changedScheduleId) broadcastSchedulerChanged(changedScheduleId);
})();
}
// 与 spend 记账并列的另一个 turn-done side-effect: 刷新 Claude 账号月度配额
// (LiteLLM /v2/user/info)。fire-and-forget, 模块内 2s 超时 + 10s 节流。
Expand Down Expand Up @@ -3456,7 +3484,23 @@ export function wireSessionToIpc(session: ReturnType<Maker['getSession']>): void
// 避免 scheduler 的 Cost 汇总混入订阅价值或远端账号消耗。
// fire-and-forget 不阻塞事件循环;价格表走 main 端内存 + 磁盘缓存,
// stale 快返并后台刷新,
// 拉不到 / 模型无条目 → 本轮不显示。
// 拉不到 / 模型无条目 → 只落 token 明细,UI 退回显示本轮 token。
// 明细在 try 外构造:它只依赖上面已拿到的 token 数,价格请求抛错时也要能落。
const turnUsageDetails = buildTurnUsageDetails({
inputTokens: promptTokens,
outputTokens: completionTokens,
cacheReadTokens: cachedTokens,
cacheCreateTokens: 0,
model: turnModel,
});
const recordCodexUsageOnly = async () => {
if (!turnAssistantPersistId) return;
await recordTurnUsageOnMessage({
sessionId: session.id,
clientId: turnAssistantPersistId,
turnUsageDetails,
});
};
try {
const pricing = isSubscriptionValue && !isCodexXaiProviderRoute
? await getModelPricing()
Expand Down Expand Up @@ -3486,13 +3530,6 @@ export function wireSessionToIpc(session: ReturnType<Maker['getSession']>): void
cacheCreateTokensDelta: 0,
});
}
const turnUsageDetails = buildTurnUsageDetails({
inputTokens: promptTokens,
outputTokens: completionTokens,
cacheReadTokens: cachedTokens,
cacheCreateTokens: 0,
model: turnModel,
});
if (money && money.amount > 0) {
if (!isSubscriptionValue) {
void recordTurnSpend(money);
Expand All @@ -3506,9 +3543,14 @@ export function wireSessionToIpc(session: ReturnType<Maker['getSession']>): void
turnOrigin: event.turnOrigin,
});
if (changedScheduleId) broadcastSchedulerChanged(changedScheduleId);
} else {
await recordCodexUsageOnly();
}
} catch {
// token row 已在价格请求前落库;价格失败只影响 API cost / message cost。
// 消息那一格仍要有事实可看:补落一次 token 明细。patch 是 agent_meta merge、
// 写的又是同一份明细,所以与上面成功分支重复执行也是幂等的(自身失败只 warn)。
await recordCodexUsageOnly();
}
})();
}
Expand Down
Loading