fix: preserve final and delayed phase content

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',

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<TMessageContentParts | undefined> = [];
@@ -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);
           }
         }

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' } },

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<string>();
+    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);