🔒 fix: Scope Detached Label Writes to Their Generation Epoch

Epoch scoping (P1). Label generation is detached and can outlive the
generation that started it. emitChunk only proves that SOME runtime is
current, not that the caller belongs to it, so an aborted generation's
fill(null) -- and its usage event -- could be attributed to whichever
generation replaced it, landing an index from the abandoned response on
top of the new one. Because an empty label renders nothing, that
overwrote content silently. emitChunk now takes an optional jobCreatedAt
and drops the event when the runtime epoch differs, mirroring the
existing setGraph/setContentParts convention, and both label emitters
pass it.

An abort now CLOSES the label scope instead of only cancelling the call:
the rejected generation still runs its catch and calls fill(null), which
would otherwise emit into a stream the next generation may already own.

Edited-response indexing (P1). The previous pass skipped the prefix
offset on resume, which was the wrong half of the problem: a sync
replaces initialResponse.content with the server's aggregatedContent,
which is completion-local, so after a reconnect its length is not the
kept-prefix length and the offset is wrong -- but it is wrong for run
steps in exactly the same way. Tool cards and the label that heads them
must share one index space; a label shifting differently from its tools
lands on another part. The label path now uses the identical expression
as useStepHandler, with no resume special-case. Correcting the
post-resume prefix length belongs in calculateContentIndex, where it
fixes both at once.

titleModel masking. The activity settings were made per-field last pass,
but the titleModel fallback a few lines below still selected an entire
config object, so a partial endpoints.all (for example one carrying only
headers) hid a named endpoint's titleModel and quietly fell the label
back to the main agent model. Both now read through one shared per-field
helper.

Resume reconciliation no longer depends solely on markActivityLabels,
which is best-effort yet had come to gate correctness: a lost flag write
silently dropped a label. The snapshot is consulted as a fallback.

The exported host type for generateLabel now admits undefined, which is
the documented "cannot serve, fall back to the direct call" signal the
hook keys on -- distinct from null, meaning it ran and produced nothing.
This commit is contained in:
Danny Avila 2026-07-26 22:40:04 -04:00
parent 46478ebed3
commit d89f8368e8
5 changed files with 95 additions and 37 deletions

View file

@ -427,10 +427,17 @@ class AgentClient extends BaseClient {
* race the persist. */
this.usageEmitSink?.push(data);
if (streamId) {
const emit = GenerationJobManager.emitChunk(streamId, {
event: UsageEvents.ON_TOKEN_USAGE,
data,
}).catch((err) => {
const emit = GenerationJobManager.emitChunk(
streamId,
{
event: UsageEvents.ON_TOKEN_USAGE,
data,
},
/** Same epoch scoping as the label event: this usage is recorded
* from a detached generation and must not bill against whichever
* generation replaced it. */
{ expectedCreatedAt: this.jobCreatedAt },
).catch((err) => {
logger.warn(`[AgentClient] Failed to emit activity-label usage: ${err?.message ?? err}`);
});
this.pendingSubagentEmits.push(emit);
@ -609,11 +616,20 @@ class AgentClient extends BaseClient {
this.activityLabelUsageSeq ??
(this.contentParts ?? []).filter((part) => part?.type === ContentTypes.ACTIVITY_LABEL).length;
this.activityLabelAbort = labelScope.abort;
/** An abort CLOSES the scope, not just cancels the call. The rejected
* generation still runs its catch and calls `fill(null)`; with the scope
* merely aborted that fill would emit and by then the next generation
* may already own the stream, so the event would land an index from the
* abandoned response onto the new one. */
const closeOnAbort = () => {
labelScope.closed = true;
labelScope.abort.abort();
};
if (abortSignal != null) {
if (abortSignal.aborted) {
labelScope.abort.abort();
closeOnAbort();
} else {
abortSignal.addEventListener('abort', () => labelScope.abort.abort(), { once: true });
abortSignal.addEventListener('abort', closeOnAbort, { once: true });
}
}
/** Thin wrapper: slot claiming, lane stamping, emit ordering, and settle
@ -640,7 +656,12 @@ class AgentClient extends BaseClient {
conversationId: this.conversationId,
},
},
{ durable: true },
/** Label generation is detached and can outlive its generation, so
* the emit is scoped to the epoch that claimed the index. Without
* it a straggler from a replaced generation lands its old index on
* the new response invisibly, since an empty label renders
* nothing overwriting whatever occupies that slot. */
{ durable: true, expectedCreatedAt: this.jobCreatedAt },
),
trackPendingFill: (fillDone) => {
this.pendingActivityLabelFills = this.pendingActivityLabelFills ?? [];

View file

@ -812,14 +812,17 @@ export default function useResumableSSE(
* is claimed in the same server-side space and needs the identical
* shift, or it lands inside the prefix and overwrites kept content.
*
* NOT on a resume: the sync replaces `initialResponse.content` with
* the server's `aggregatedContent`, which already contains the prefix
* AND everything generated since. Its length is not the prefix
* length, and the indices reconciled from that snapshot are already
* absolute offsetting again would push the label past its slot and
* overwrite a later part. */
* Deliberately the SAME expression `useStepHandler` uses, with no
* resume special-case. A sync replaces `initialResponse.content` with
* the server's `aggregatedContent`, which is completion-local so
* after a reconnect its length is not the kept-prefix length and this
* offset is wrong. It is wrong for run steps in exactly the same way,
* and tool cards and their label MUST share one index space: a label
* that shifts differently from the tools it heads would land on
* another part. Fixing the post-resume prefix length belongs in
* `calculateContentIndex`, where it corrects both at once. */
const initialContent =
!isResume && currentSubmission.editedContent != null
currentSubmission.editedContent != null
? ((currentSubmission.initialResponse as TMessage | undefined)?.content ?? [])
: [];
const offsetEvent =

View file

@ -91,25 +91,37 @@ export interface ResolvedActivityConfig {
* theirs: an `endpoints.all` block wins over the named endpoint, which wins
* over a custom endpoint's own config.
*/
/**
* Reads ONE endpoint setting, global-then-named, rather than picking a whole
* config object.
*
* Selecting wholesale means any `endpoints.all` block even one carrying
* nothing but `headers` shadows the named/custom endpoint entirely, so a
* single unrelated global setting silently hides every activity field AND the
* `titleModel` fallback. Global still wins per field, so a real
* `all.activityLabel` keeps overriding the endpoint.
*/
function pickEndpointField<K extends keyof TEndpoint>(
appConfig: AppConfig | undefined,
endpoint: string,
customEndpointConfig: Partial<TEndpoint> | undefined,
key: K,
): TEndpoint[K] | undefined {
const endpoints = appConfig?.endpoints as
| (Record<string, TEndpoint | undefined> & { all?: TEndpoint })
| undefined;
const all = endpoints?.all as Partial<TEndpoint> | undefined;
const named = (endpoints?.[endpoint] ?? customEndpointConfig) as Partial<TEndpoint> | undefined;
return all?.[key] ?? named?.[key];
}
export function resolveActivityConfig(
appConfig: AppConfig | undefined,
endpoint: string,
customEndpointConfig?: Partial<TEndpoint>,
): ResolvedActivityConfig {
const endpoints = appConfig?.endpoints as
| (Record<string, TEndpoint | undefined> & { all?: TEndpoint })
| undefined;
/**
* Resolved FIELD BY FIELD rather than by picking one config object whole.
* Selecting wholesale means any `endpoints.all` block even one carrying
* nothing but `headers` shadows the named/custom endpoint entirely and
* silently disables activity labels everywhere. Global still wins per
* field, so a real `all.activityLabel` keeps overriding the endpoint.
*/
const all = endpoints?.all as Partial<TEndpoint> | undefined;
const named = (endpoints?.[endpoint] ?? customEndpointConfig) as Partial<TEndpoint> | undefined;
const pick = <K extends keyof TEndpoint>(key: K): TEndpoint[K] | undefined =>
all?.[key] ?? named?.[key];
pickEndpointField(appConfig, endpoint, customEndpointConfig, key);
return {
enabled: pick('activityLabel') === true,
model: pick('activityModel'),
@ -159,15 +171,19 @@ export async function resolveActivityLabelModel({
}
}
const endpoints = appConfig?.endpoints as
| (Record<string, TEndpoint | undefined> & { all?: TEndpoint })
| undefined;
const endpointConfig: Partial<TEndpoint> | undefined =
endpoints?.all ?? endpoints?.[endpoint] ?? providerConfig.customEndpointConfig;
/** Same per-field read as the activity settings: a partial `endpoints.all`
* must not hide the resolved endpoint's `titleModel` and quietly fall the
* label back to the main agent's (usually much larger) model. */
const titleModel = pickEndpointField(
appConfig,
endpoint,
providerConfig.customEndpointConfig,
'titleModel',
);
const model =
activity.model ??
(endpointConfig?.titleModel != null && endpointConfig.titleModel !== Constants.CURRENT_MODEL
? endpointConfig.titleModel
(titleModel != null && titleModel !== Constants.CURRENT_MODEL
? titleModel
: (agent.model ?? agent.model_parameters?.model));
const options = await providerConfig.getOptions({
req,

View file

@ -188,7 +188,13 @@ export interface ActivityLabelHostDeps {
*/
isClosed?: () => boolean;
resolveLLM: () => Promise<ActivityLabelLLM>;
generateLabel?: (payload: GenerateLabelPayload) => Promise<string | null>;
/**
* Resolve `undefined` to DECLINE this bridge cannot serve the request, so
* the hook falls back to the direct model call. `null` means it ran and
* produced no label. The distinction is the contract the hook keys on, so it
* belongs in the exported type.
*/
generateLabel?: (payload: GenerateLabelPayload) => Promise<string | null | undefined>;
getInvokeCallbacks?: () => ActivityLabelInvokeCallbacks;
}

View file

@ -2297,7 +2297,20 @@ class GenerationJobManagerClass {
// Compare the snapshot content view against a fresh read and re-emit any
// label whose text/pending state moved; the client applier is idempotent
// and refuses stale pending placeholders.
if (resumeState != null && jobActive && liveJob?.activityLabels === true) {
/** The flag is the fast path, but `markActivityLabels` is best-effort
* and now gates correctness rather than merely saving a read a lost
* write would silently drop a label. Fall back to the snapshot when it
* is absent: that misses only the case where the FIRST label is claimed
* inside the gap, which the flag covers whenever it did persist. */
const snapshotHasActivityLabels =
resumeState?.aggregatedContent?.some(
(part) => (part as { type?: string } | null)?.type === 'activity_label',
) === true;
if (
resumeState != null &&
jobActive &&
(liveJob?.activityLabels === true || snapshotHasActivityLabels)
) {
const labelContent = await this.jobStore.getContentParts(streamId, liveJob.createdAt);
if (options?.signal?.aborted || this.detachSubscriptionDuringShutdown(subscription)) {
return cancelResumeSubscription();
@ -2408,7 +2421,6 @@ class GenerationJobManagerClass {
) {
return;
}
const sequence = ++runtime.emissionSequence;
let signalSnapshotReady!: () => void;
const snapshotReady = new Promise<void>((resolve) => {