🧾 fix: Keep Label Accounting Out of the Primary Usage Slot

Label usage no longer owns getStreamUsage(). recordCollectedUsage assigned
its result to this.usage unconditionally, so when the primary provider
reported no usage metadata but the label provider did, BaseClient took the
label's output tokens as the assistant response's authoritative count. The
later primary call returns early on an empty collectedUsage and never
replaced it, so the wrong value stood, the text-based fallback was skipped,
and the real generation went unbilled. Secondary usage is still billed but
no longer writes that slot.

Cross-endpoint pricing keys off an explicit discriminator rather than the
presence of a value. A built-in label endpoint prices from the shared
table, so an undefined endpointTokenConfig is its MEANINGFUL value --
reading that as "no override" fell back to the primary's custom rates and
restored the exact mismatch the previous pass set out to fix. The caller
already knows whether the label ran elsewhere and now says so.

markActivityLabels rejects on failure instead of swallowing it. The flag
gates resume gap reconciliation and the caller retries it, but the internal
catch resolved successfully and made that retry unreachable -- so the two
changes cancelled out and a transient write failure still left the flag
absent.

Late label accounting is suppressed with the same gate as the late fill. A
straggler that outlived the settle timeout still ran its finally block, so
it charged the balance and appended to usageEmitSink after the response had
passed its usage flush and metadata snapshot: a cost the user pays but is
never shown.

The cleared-prefix state is scoped to one generation. It was set on a
resume SYNC that replaced the response and then never reset, so a later
edited resubmission in the same mounted hook dispatched run steps and
labels with no offset against content that still held its retained prefix.
Reconnects pass isResume and keep the state; a new generation clears it.
This commit is contained in:
Danny Avila 2026-07-27 09:01:58 -04:00
parent 57e45cbb67
commit d8876a194f
3 changed files with 64 additions and 18 deletions

View file

@ -462,9 +462,13 @@ class AgentClient extends BaseClient {
context: 'activity-label',
model,
endpointTokenConfig: labelTokenConfig,
/** The label ran elsewhere, so its config governs even when undefined. */
crossEndpoint: sameEndpoint === false,
balance: getBalanceConfig(appConfig),
transactions: getTransactionsConfig(appConfig),
messageId: this.responseMessageId,
/** Billed, but NOT the response's stream usage — see the parameter. */
updateStreamUsage: false,
}).catch((err) => {
logger.error(
'[api/server/controllers/agents/client.js #recordActivityLabelUsage] Error recording usage',
@ -537,12 +541,25 @@ class AgentClient extends BaseClient {
});
return label ?? null;
} finally {
await this.recordActivityLabelUsage(
collectedMetadata,
clientOptions.model,
endpointTokenConfig,
sameEndpoint,
);
/**
* Skip accounting for a straggler that outlived its scope. Settle has
* already timed out, the response is past its usage flush and metadata
* snapshot, and the fill itself is dropped so recording here would
* charge the balance and push into `usageEmitSink` after anything can
* surface it, producing a cost the user is billed for but never shown.
* Late bookkeeping is suppressed with the same gate as the late fill.
*/
const allScopesClosed =
(this.activityLabelScopes ?? []).length > 0 &&
(this.activityLabelScopes ?? []).every((scope) => scope.closed === true);
if (!allScopesClosed) {
await this.recordActivityLabelUsage(
collectedMetadata,
clientOptions.model,
endpointTokenConfig,
sameEndpoint,
);
}
}
}
@ -1639,11 +1656,33 @@ class AgentClient extends BaseClient {
* the label's, so the two disagreed. `undefined` keeps the agent default.
*/
endpointTokenConfig,
/**
* True when this usage ran on a DIFFERENT endpoint than the agent, making
* `endpointTokenConfig` authoritative even when it is `undefined` (a
* built-in endpoint prices from the shared table). Presence of the value
* cannot express that, which is why the caller states it outright.
*/
crossEndpoint = false,
/**
* Whether this recording owns `getStreamUsage()`. Only the PRIMARY
* generation does. Secondary usage (activity labels) must still be
* billed, but writing it here would hand `BaseClient` the label's token
* counts as the assistant response's authoritative total and because
* the primary call returns early when it collected nothing, the wrong
* value would never be replaced, suppressing the text-based token
* fallback and leaving the real generation unbilled.
*/
updateStreamUsage = true,
}) {
/** Per-agent resolution keys off the AGENT's config map, which cannot
* describe a label running on a different endpoint so an explicit
* config wins outright rather than being second-guessed per usage row. */
const overrideTokenConfig = endpointTokenConfig !== undefined;
* config wins outright rather than being second-guessed per usage row.
*
* Keyed on the caller's discriminator, NOT on `endpointTokenConfig !==
* undefined`: a built-in label endpoint prices from the shared table, so
* `undefined` is its meaningful value. Reading that as "no override" is
* what silently restored the primary's custom rates. */
const overrideTokenConfig = crossEndpoint === true;
const result = await recordCollectedUsage(
{
spendTokens: db.spendTokens,
@ -1669,7 +1708,7 @@ class AgentClient extends BaseClient {
},
);
if (result) {
if (result && updateStreamUsage) {
this.usage = result;
}
}

View file

@ -695,6 +695,17 @@ export default function useResumableSSE(
*/
const subscribeToStream = useCallback(
(currentStreamId: string, currentSubmission: TSubmission, isResume = false) => {
/**
* A NEW generation starts with its retained prefix intact, so the
* cleared-prefix state from a previous one must not carry over the
* hook outlives any single submission, and a later edited resubmission
* in the same chat would otherwise be dispatched with no offset and
* overwrite the content it kept. Reconnects pass `isResume`, so the
* state survives exactly where it should: within one generation.
*/
if (!isResume) {
editPrefixClearedRef.current = false;
}
let { userMessage } = currentSubmission;
let textIndex: number | null = null;
let finalReceived = false;

View file

@ -2421,15 +2421,11 @@ class GenerationJobManagerClass {
* Best-effort: the flag is an optimization hint, never correctness.
*/
async markActivityLabels(streamId: string): Promise<void> {
try {
await this.jobStore.updateJob(streamId, { activityLabels: true });
} catch (error) {
logger.debug(
`[GenerationJobManager] Could not flag activity labels for ${streamId}: ${
(error as Error)?.message ?? error
}`,
);
}
/** Deliberately REJECTS on failure. This flag gates resume gap
* reconciliation, so the caller retries it; swallowing the error here
* resolved successfully and made that retry unreachable, leaving the
* flag absent after a transient write failure. */
await this.jobStore.updateJob(streamId, { activityLabels: true });
}
async emitChunk(