From 59594ec17100db8cc2cebd729c3cb9495d05b4f8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 22 Jul 2026 12:02:49 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A6=20fix:=20Codex=20Round=206=20?= =?UTF-8?q?=E2=80=94=20Stream=20Label=20Usage,=20Close=20Late=20Fills?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Emit an on_token_usage chunk for label calls (sink push alone left the live session gauge blind); retained in pendingSubagentEmits so job cleanup cannot race the persist, tagged 'activity-label' as before. - Close the label scope when settle times out: the wiring gates fill() on isClosed and the client fires a label-scoped AbortController, so a straggling generation can neither mutate a saved response nor emit into a job whose runtime is gone. The controller also chains to the run signal, so a user abort still cancels label work. --- api/server/controllers/agents/client.js | 45 ++++++++++++++++--- .../activityLabels/__tests__/wiring.spec.ts | 44 ++++++++++++++++++ .../api/src/agents/activityLabels/host.ts | 12 ++++- .../api/src/agents/activityLabels/wiring.ts | 14 ++++++ 4 files changed, 109 insertions(+), 6 deletions(-) diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index db3b862bca..d7ccf2d3f6 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -374,15 +374,30 @@ class AgentClient extends BaseClient { if (collectedUsage.length === 0) { return; } + const streamId = this.options.req?._resumableStreamId || null; for (const usage of collectedUsage) { - this.usageEmitSink?.push({ + const data = { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, model, usage_type: 'activity-label', runId: this.responseMessageId, seq: this.collectedUsage.length, - }); + }; + /** Fold into the response rollup synchronously, then stream it like + * primary/subagent usage so the live session gauge stays honest. + * Retained and flushed with the subagent emits so job cleanup cannot + * race the persist. */ + this.usageEmitSink?.push(data); + if (streamId) { + const emit = GenerationJobManager.emitChunk(streamId, { + event: UsageEvents.ON_TOKEN_USAGE, + data, + }).catch((err) => { + logger.warn(`[AgentClient] Failed to emit activity-label usage: ${err?.message ?? err}`); + }); + this.pendingSubagentEmits.push(emit); + } } await this.recordCollectedUsage({ collectedUsage, @@ -441,14 +456,19 @@ class AgentClient extends BaseClient { } } - /** Bounded settle for in-flight label fills before finalization. */ + /** Bounded settle for in-flight label fills before finalization. On + * timeout the label scope is closed and its abort controller fired, so a + * straggler cannot mutate the saved response or emit into a dead job. */ async settleActivityLabels(timeoutMs = 3000) { const pending = this.pendingActivityLabelFills; if (!pending || pending.length === 0) { return; } this.pendingActivityLabelFills = []; - await settlePendingLabelFills(pending, timeoutMs); + await settlePendingLabelFills(pending, timeoutMs, () => { + this.activityLabelsClosed = true; + this.activityLabelAbort?.abort(); + }); } /** @@ -467,10 +487,25 @@ class AgentClient extends BaseClient { /** SDK support probe (steering-style): the Run method and the formatter * replay skip ship together, so method presence is the capability. */ const sdkCapable = typeof Run?.prototype?.generateActivityLabel === 'function'; + /** Label-scoped abort: fired when settle times out so a straggling + * generation stops burning provider time for a finalized response. + * Chained to the run signal so a user abort still cancels labels. */ + this.activityLabelsClosed = false; + this.activityLabelAbort = new AbortController(); + if (abortSignal != null) { + if (abortSignal.aborted) { + this.activityLabelAbort.abort(); + } else { + abortSignal.addEventListener('abort', () => this.activityLabelAbort?.abort(), { + once: true, + }); + } + } /** Thin wrapper: slot claiming, lane stamping, emit ordering, and settle * tracking live in `createActivityLabelWiring` (packages/api, TS). */ return createActivityLabelWiring({ - abortSignal, + abortSignal: this.activityLabelAbort.signal, + isClosed: () => this.activityLabelsClosed === true, getContentParts: () => this.contentParts, bumpIndexOffset: () => { this.steerOffsetState.offset += 1; diff --git a/packages/api/src/agents/activityLabels/__tests__/wiring.spec.ts b/packages/api/src/agents/activityLabels/__tests__/wiring.spec.ts index d1d9f6c6f3..40220e5687 100644 --- a/packages/api/src/agents/activityLabels/__tests__/wiring.spec.ts +++ b/packages/api/src/agents/activityLabels/__tests__/wiring.spec.ts @@ -141,3 +141,47 @@ describe('synthesizeActivityLabelGapEvents', () => { expect(synthesizeActivityLabelGapEvents(parts, parts, meta)).toEqual([]); }); }); + +describe('createActivityLabelWiring close gate', () => { + it('drops a late fill once the response has finalized', async () => { + const parts: Array = [ + { type: 'tool_call', tool_call: { id: 'tool-1' } }, + ]; + const emitLabelEvent = jest.fn(async () => undefined); + let closed = false; + let releaseLabel: (value: string) => void = () => undefined; + const generateLabel = jest.fn( + () => + new Promise((resolve) => { + releaseLabel = resolve; + }), + ); + const { hook } = createActivityLabelWiring({ + getContentParts: () => parts, + bumpIndexOffset: jest.fn(), + emitLabelEvent, + trackPendingFill: jest.fn(), + isClosed: () => closed, + resolveLLM: jest.fn(async () => ({ + provider: Providers.OPENAI, + clientOptions: { model: 'm' }, + })), + generateLabel, + }); + + await hook(batchInput(), new AbortController().signal); + await flushDetached(); + /** Claim-time placeholder emitted; label still in flight. */ + expect(emitLabelEvent).toHaveBeenCalledTimes(1); + + /** Settle timed out: the scope closes, then the straggler resolves. */ + closed = true; + releaseLabel('Late label that must not land'); + await flushDetached(); + + expect(emitLabelEvent).toHaveBeenCalledTimes(1); + const labelPart = parts[1] as LooseContentPart; + expect(labelPart.activity_label).toBe(''); + expect(labelPart.pending).toBe(true); + }); +}); diff --git a/packages/api/src/agents/activityLabels/host.ts b/packages/api/src/agents/activityLabels/host.ts index e66b5c318e..5d8272ccb6 100644 --- a/packages/api/src/agents/activityLabels/host.ts +++ b/packages/api/src/agents/activityLabels/host.ts @@ -138,16 +138,26 @@ export async function resolveActivityLabelModel({ export async function settlePendingLabelFills( pending: Array>, timeoutMs = 3000, + onTimeout?: () => void, ): Promise { if (pending.length === 0) { return; } let timerId: ReturnType | undefined; + let timedOut = false; const timeout = new Promise((resolve) => { - timerId = setTimeout(resolve, timeoutMs); + timerId = setTimeout(() => { + timedOut = true; + resolve(); + }, timeoutMs); }); await Promise.race([Promise.allSettled(pending), timeout]); if (timerId != null) { clearTimeout(timerId); } + if (timedOut) { + /** Stragglers must not mutate or emit for a response that is already + * finalizing: the caller aborts them and closes the slot gate. */ + onTimeout?.(); + } } diff --git a/packages/api/src/agents/activityLabels/wiring.ts b/packages/api/src/agents/activityLabels/wiring.ts index 26de7e7db7..a2e0f19f58 100644 --- a/packages/api/src/agents/activityLabels/wiring.ts +++ b/packages/api/src/agents/activityLabels/wiring.ts @@ -175,6 +175,12 @@ export interface ActivityLabelHostDeps { emitLabelEvent: (index: number, part: LooseContentPart) => Promise; /** Registers a fill-completion promise for bounded settle at finalization. */ trackPendingFill: (fillDone: Promise) => void; + /** + * True once the response has finalized (settle timed out). A late fill + * must then neither mutate persisted content nor emit chunks for a job + * whose runtime is gone. + */ + isClosed?: () => boolean; resolveLLM: () => Promise; generateLabel?: (payload: GenerateLabelPayload) => Promise; getInvokeCallbacks?: () => ActivityLabelInvokeCallbacks; @@ -246,12 +252,20 @@ export function createActivityLabelWiring(deps: ActivityLabelHostDeps): { context, fill: async (text) => { try { + /** Finalization already passed: drop the result rather than + * mutating a saved response or emitting into a closed job. */ + if (deps.isClosed?.() === true) { + return; + } part.pending = false; if (text == null || text.length === 0) { return; } part[ContentTypes.ACTIVITY_LABEL] = text; await claimEmit; + if (deps.isClosed?.() === true) { + return; + } await deps.emitLabelEvent(index, part); } finally { resolveFill();