From 831a0035302159220b5fd72a4547c672ec4af678 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 13 Aug 2026 09:23:05 -0400 Subject: [PATCH] fix: preserve final and delayed phase content --- .../utils/__tests__/activityLabels.spec.ts | 58 ++++++++++ client/src/utils/activityLabels.ts | 10 +- .../src/agents/activityPhases/runtime.spec.ts | 105 ++++++++++++++++++ .../api/src/agents/activityPhases/runtime.ts | 40 +++++-- 4 files changed, 197 insertions(+), 16 deletions(-) diff --git a/client/src/utils/__tests__/activityLabels.spec.ts b/client/src/utils/__tests__/activityLabels.spec.ts index fa9fadbc3f..9493822e6f 100644 --- a/client/src/utils/__tests__/activityLabels.spec.ts +++ b/client/src/utils/__tests__/activityLabels.spec.ts @@ -284,6 +284,64 @@ describe('groupActivityPhases', () => { ]); }); + it('moves a late child label from an earlier phase into its declared later span', () => { + const firstTool = { + type: ContentTypes.TOOL_CALL, + tool_call: { id: 'first', name: 'web_search', args: '{}', output: 'ok' }, + } as unknown as TMessageContentParts; + const laterTool = { + type: ContentTypes.TOOL_CALL, + tool_call: { id: 'later', name: 'web_search', args: '{}', output: 'ok' }, + } as unknown as TMessageContentParts; + const lateChild = labelPart({ + activity_label: 'Recorded the later result', + pending: false, + tool_call_ids: ['later'], + }); + const first = labelPart({ activity_label: 'Completed the first phase', pending: false }); + Object.assign(first, { + activity_label_type: 'phase', + activity_start_index: 0, + activity_end_index: 1, + activity_count: 2, + }); + const second = labelPart({ activity_label: 'Completed the second phase', pending: false }); + Object.assign(second, { + activity_label_type: 'phase', + activity_start_index: 2, + activity_end_index: 4, + activity_count: 2, + }); + const boundary = { type: ContentTypes.TEXT, text: 'Boundary' } as TMessageContentParts; + const content = [ + firstTool, + boundary, + laterTool, + lateChild as never, + first as never, + second as never, + ]; + + const segments = groupActivityPhases(content); + + expect(segments).toEqual([ + expect.objectContaining({ + type: 'phase', + labelIndex: 4, + contentIndices: [0], + }), + expect.objectContaining({ + type: 'content', + contentIndices: [1], + }), + expect.objectContaining({ + type: 'phase', + labelIndex: 5, + contentIndices: [2, 3], + }), + ]); + }); + it('keeps a late child label in the phase while leaving final text outside', () => { const child = labelPart({ activity_label: 'Recorded the delayed child result', diff --git a/client/src/utils/activityLabels.ts b/client/src/utils/activityLabels.ts index 736e33b779..fcfb7fcbc6 100644 --- a/client/src/utils/activityLabels.ts +++ b/client/src/utils/activityLabels.ts @@ -195,7 +195,7 @@ export function groupActivityPhases( const deferredTrailingIndices: number[] = []; for (let segmentIndex = segments.length - 1; segmentIndex >= 0; segmentIndex -= 1) { const segment = segments[segmentIndex]; - if (segment.type !== 'content') { + if (segment.type !== 'content' && segment.type !== 'phase') { continue; } const retainedContent: Array = []; @@ -206,12 +206,14 @@ export function groupActivityPhases( childPosition += 1 ) { const childIndex = segment.contentIndices[childPosition]; - if (childIndex >= start && childIndex < end) { + const child = segment.content[childPosition]; + const canRecover = segment.type === 'content' || getBatchActivityLabelPart(child) != null; + if (canRecover && childIndex >= start && childIndex < end) { recoveredIndices.push(childIndex); - } else if (childIndex >= end) { + } else if (canRecover && childIndex >= end) { deferredTrailingIndices.push(childIndex); } else { - retainedContent.push(segment.content[childPosition]); + retainedContent.push(child); retainedIndices.push(childIndex); } } diff --git a/packages/api/src/agents/activityPhases/runtime.spec.ts b/packages/api/src/agents/activityPhases/runtime.spec.ts index 735e816aea..15e9650072 100644 --- a/packages/api/src/agents/activityPhases/runtime.spec.ts +++ b/packages/api/src/agents/activityPhases/runtime.spec.ts @@ -2675,6 +2675,111 @@ describe('createActivityPhaseWiring', () => { expect(wiring.snapshot().activityCount).toBe(0); }); + it('retains uncovered calls from a delayed straddling tool batch', async () => { + const parts: LooseContentPart[] = [ + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-1' } }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-2' } }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'covered-call' } }, + { type: ContentTypes.TEXT, text: substantialText('Boundary result.') }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'uncovered-call' } }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-3' } }, + ]; + const payloads: GenerateActivityPhasePayload[] = []; + const wiring = createActivityPhaseWiring({ + getContentParts: () => parts, + bumpIndexOffset: jest.fn(), + emitLabelEvent: jest.fn(async () => undefined), + trackPendingFill: jest.fn(), + generatePhase: jest.fn(async (payload: GenerateActivityPhasePayload) => { + payloads.push(payload); + return { label: 'Completed one phase' }; + }), + }); + await wiring.hook(batch('tool-1'), new AbortController().signal); + await wiring.hook(batch('tool-2'), new AbortController().signal); + wiring.complete(); + + const delayedBatch = batch('covered-call'); + delayedBatch.entries.push({ + ...delayedBatch.entries[0], + toolUseId: 'uncovered-call', + toolInput: { query: 'uncovered-call' }, + }); + await wiring.hook(delayedBatch, new AbortController().signal); + await wiring.hook(batch('tool-3'), new AbortController().signal); + wiring.complete(); + await flushDetached(); + + expect(payloads.map(({ totalActivityCount }) => totalActivityCount)).toEqual([2, 2]); + expect(payloads[1].activities[0]).toMatchObject({ + entries: [expect.objectContaining({ toolInput: { query: 'uncovered-call' } })], + }); + }); + + it('keeps a short semantic final answer outside the collapsed phase', async () => { + const parts: LooseContentPart[] = [ + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-1' } }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-2' } }, + ]; + const wiring = createActivityPhaseWiring({ + getContentParts: () => parts, + getStepIndex: (stepId) => (stepId === 'final-step' ? 2 : undefined), + bumpIndexOffset: jest.fn(), + emitLabelEvent: jest.fn(async () => undefined), + trackPendingFill: jest.fn(), + generatePhase: jest.fn(async () => ({ label: 'Completed the requested work' })), + }); + await wiring.hook(batch('tool-1'), new AbortController().signal); + await wiring.hook(batch('tool-2'), new AbortController().signal); + const handlers = wiring.handlers({ + [GraphEvents.ON_RUN_STEP]: { + handle: () => { + parts[2] = { type: ContentTypes.TEXT, text: '', phase: 'final_answer' }; + }, + }, + [GraphEvents.ON_MESSAGE_DELTA]: { + handle: (_event, data) => { + const delta = data as { delta?: { content?: { text?: string } } }; + parts[2] = { + type: ContentTypes.TEXT, + text: `${parts[2]?.text ?? ''}${delta.delta?.content?.text ?? ''}`, + phase: 'final_answer', + }; + }, + }, + }); + handlers?.[GraphEvents.ON_RUN_STEP]?.handle( + GraphEvents.ON_RUN_STEP, + { + id: 'final-step', + stepDetails: { + type: StepTypes.MESSAGE_CREATION, + message_creation: { message_id: 'm', content_type: 'text', phase: 'final_answer' }, + }, + }, + undefined, + undefined, + ); + handlers?.[GraphEvents.ON_MESSAGE_DELTA]?.handle( + GraphEvents.ON_MESSAGE_DELTA, + { + id: 'final-step', + delta: { content: { type: ContentTypes.TEXT, text: 'Done.' } }, + }, + undefined, + undefined, + ); + await flushDetached(); + + expect(parts[2]).toMatchObject({ type: ContentTypes.TEXT, text: 'Done.' }); + expect(parts[3]).toMatchObject({ + activity_label_type: 'phase', + activity_start_index: 0, + activity_end_index: 2, + activity_count: 2, + }); + }); + it('retains equal-position context rendered after a substantial boundary', async () => { const parts: LooseContentPart[] = [ { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-1' } }, diff --git a/packages/api/src/agents/activityPhases/runtime.ts b/packages/api/src/agents/activityPhases/runtime.ts index 71b53273d0..7030e77bdc 100644 --- a/packages/api/src/agents/activityPhases/runtime.ts +++ b/packages/api/src/agents/activityPhases/runtime.ts @@ -1177,7 +1177,8 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity * subagent internals. */ const reasoningKey = input.executingAgentId ?? 'root'; const reasoning = pendingReasoning.get(reasoningKey)?.text.trim(); - const ids = new Set(input.entries.map((entry) => entry.toolUseId)); + let trackedEntries = input.entries; + let ids = new Set(trackedEntries.map((entry) => entry.toolUseId)); const parts = deps.getContentParts(); let childLabelIndex: number | undefined; let batchStartIndex: number | undefined; @@ -1203,23 +1204,27 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity } } } - const batchToolIndices = definedPartIndices(parts).filter((index) => { + const coveredToolIds = new Set(); + for (const index of definedPartIndices(parts)) { const part = parts[index]; - return ( + if ( part?.type === ContentTypes.TOOL_CALL && typeof part.tool_call?.id === 'string' && - ids.has(part.tool_call.id) - ); - }); - if ( - batchToolIndices.some((toolIndex) => - emittedContentRanges.some(({ start, end }) => toolIndex >= start && toolIndex < end), - ) - ) { + ids.has(part.tool_call.id) && + emittedContentRanges.some(({ start, end }) => index >= start && index < end) + ) { + coveredToolIds.add(part.tool_call.id); + } + } + if (coveredToolIds.size > 0) { + trackedEntries = trackedEntries.filter((entry) => !coveredToolIds.has(entry.toolUseId)); + ids = new Set(trackedEntries.map((entry) => entry.toolUseId)); + } + if (trackedEntries.length === 0) { pendingReasoning.delete(reasoningKey); return {}; } - const entries = input.entries.map((entry: BatchEntry) => ({ + const entries = trackedEntries.map((entry: BatchEntry) => ({ toolName: entry.toolName, toolInput: entry.toolInput, toolOutput: entry.toolOutput, @@ -1311,6 +1316,17 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity textContextByStepId.delete(removed.stepId); } } + const result = runStepHandler.handle(event, data, metadata, graph); + if (phase === 'final_answer' && isRoot) { + const boundaryIndex = deps.getStepIndex?.(step.id); + if ( + boundaryIndex != null && + deps.getContentParts()[boundaryIndex]?.type === ContentTypes.TEXT + ) { + close(phase, boundaryIndex); + } + } + return result; } } return runStepHandler.handle(event, data, metadata, graph);