diff --git a/packages/api/src/agents/activityPhases/runtime.spec.ts b/packages/api/src/agents/activityPhases/runtime.spec.ts index 49307fb53f..735e816aea 100644 --- a/packages/api/src/agents/activityPhases/runtime.spec.ts +++ b/packages/api/src/agents/activityPhases/runtime.spec.ts @@ -1307,7 +1307,7 @@ describe('createActivityPhaseWiring', () => { expect(parts[3]).toBeUndefined(); }); - it('bounds persisted evidence while preserving the full activity count', async () => { + it('bounds persisted evidence and the tracked activity window', async () => { const parts: LooseContentPart[] = []; const wiring = createActivityPhaseWiring({ getContentParts: () => parts, @@ -1324,7 +1324,7 @@ describe('createActivityPhaseWiring', () => { const snapshot = wiring.snapshot(); expect(snapshot.version).toBe(2); - expect(snapshot.activityCount).toBe(100); + expect(snapshot.activityCount).toBe(77); expect(snapshot.activities).toHaveLength(13); expect(snapshot.overflowActivityStartIndex).toBe(99); expect(snapshot.overflowToolCallIds).toHaveLength(64); @@ -2758,6 +2758,76 @@ describe('createActivityPhaseWiring', () => { expect(payloads[0].assistantContext).toBeUndefined(); expect(payloads[1].assistantContext).toEqual(['Context for the later work.']); }); + + it('retains a completed tool batch when any call crosses a substantial boundary', 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: 'parallel-1' } }, + { type: ContentTypes.TEXT, text: '' }, + { type: ContentTypes.TOOL_CALL, tool_call: { id: 'parallel-2' } }, + ]; + const payloads: GenerateActivityPhasePayload[] = []; + const wiring = createActivityPhaseWiring({ + getContentParts: () => parts, + getStepIndex: (stepId) => (stepId === 'boundary' ? 3 : undefined), + 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); + const parallelBatch = batch('parallel-1'); + parallelBatch.entries.push({ + ...parallelBatch.entries[0], + toolUseId: 'parallel-2', + toolInput: { query: 'parallel-2' }, + }); + await wiring.hook(parallelBatch, new AbortController().signal); + const handlers = wiring.handlers({ + [GraphEvents.ON_RUN_STEP]: { handle: jest.fn() }, + [GraphEvents.ON_MESSAGE_DELTA]: { + handle: (_event, data) => { + const delta = data as { delta?: { content?: { text?: string } } }; + parts[3] = { type: ContentTypes.TEXT, text: delta.delta?.content?.text ?? '' }; + }, + }, + }); + handlers?.[GraphEvents.ON_RUN_STEP]?.handle( + GraphEvents.ON_RUN_STEP, + { + id: 'boundary', + stepDetails: { + type: StepTypes.MESSAGE_CREATION, + message_creation: { message_id: 'm', content_type: 'text' }, + }, + }, + undefined, + undefined, + ); + handlers?.[GraphEvents.ON_MESSAGE_DELTA]?.handle( + GraphEvents.ON_MESSAGE_DELTA, + { + id: 'boundary', + delta: { content: { type: ContentTypes.TEXT, text: substantialText('Boundary result.') } }, + }, + undefined, + undefined, + ); + parts[6] = { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-3' } }; + await wiring.hook(batch('tool-3'), new AbortController().signal); + + wiring.complete(); + await flushDetached(); + + expect(payloads.map(({ totalActivityCount }) => totalActivityCount)).toEqual([2, 2]); + expect(parts[5]).toMatchObject({ activity_end_index: 3, activity_count: 2 }); + expect(parts[7]).toMatchObject({ activity_start_index: 4, activity_count: 2 }); + }); }); describe('createAssistantPhaseStampingHandlers', () => { diff --git a/packages/api/src/agents/activityPhases/runtime.ts b/packages/api/src/agents/activityPhases/runtime.ts index 9627bfab6b..71b53273d0 100644 --- a/packages/api/src/agents/activityPhases/runtime.ts +++ b/packages/api/src/agents/activityPhases/runtime.ts @@ -26,6 +26,8 @@ export interface ActivityPhaseEntry { export type TrackedActivity = ActivityPhaseEntry & { startIndex: number; + /** A prior boundary can retain only the materialized tail of a straddling batch. */ + partitionStartIndex?: number; childLabelIndex?: number; /** Stable anchors survive content filtering and prepends across HITL resume. */ toolCallIds?: string[]; @@ -216,30 +218,14 @@ function findLastPartIndex( return Math.max(0, parts.length - 1); } -function findBatchStart( - parts: ReadonlyArray, - toolCallIds: Set, -): number { - let first = -1; - for (const index of definedPartIndices(parts)) { - const part = parts[index]; - if ( - part?.type === ContentTypes.TOOL_CALL && - typeof part.tool_call?.id === 'string' && - toolCallIds.has(part.tool_call.id) - ) { - first = first < 0 ? index : Math.min(first, index); - } - } - return first >= 0 ? first : Math.max(0, parts.length - 1); -} - function findTrackedStart( parts: ReadonlyArray, activity: TrackedActivity, ): number { const materializedStart = findMaterializedActivityStart(parts, activity); - return materializedStart ?? Math.min(activity.startIndex, Math.max(0, parts.length - 1)); + const startIndex = + materializedStart ?? Math.min(activity.startIndex, Math.max(0, parts.length - 1)); + return Math.max(startIndex, activity.partitionStartIndex ?? 0); } function findMaterializedActivityStart( @@ -312,16 +298,25 @@ function findTrackedToolStart( parts: ReadonlyArray, activity: TrackedActivity, ): number | undefined { - if (activity.toolCallIds != null && activity.toolCallIds.length > 0) { - const toolStart = findBatchStart(parts, new Set(activity.toolCallIds)); - if ( - parts[toolStart]?.type === ContentTypes.TOOL_CALL && - activity.toolCallIds.includes(String(parts[toolStart]?.tool_call?.id ?? '')) - ) { - return toolStart; - } + return findTrackedToolIndices(parts, activity)[0]; +} + +function findTrackedToolIndices( + parts: ReadonlyArray, + activity: TrackedActivity, +): number[] { + if (activity.toolCallIds == null || activity.toolCallIds.length === 0) { + return []; } - return undefined; + const ids = new Set(activity.toolCallIds); + return definedPartIndices(parts).filter((index) => { + const part = parts[index]; + return ( + part?.type === ContentTypes.TOOL_CALL && + typeof part.tool_call?.id === 'string' && + ids.has(part.tool_call.id) + ); + }); } function findReasoningStart( @@ -715,6 +710,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity } else { overflowActivities.push({ startIndex: activity.startIndex, + ...(activity.partitionStartIndex != null && { + partitionStartIndex: activity.partitionStartIndex, + }), ...(activity.toolCallIds != null && { toolCallIds: activity.toolCallIds.slice(-MAX_RETAINED_TOOL_ENTRIES), }), @@ -727,7 +725,13 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity status: activity.status, }); if (overflowActivities.length > MAX_OVERFLOW_ACTIVITY_ANCHORS) { - overflowActivities.shift(); + const removed = overflowActivities.shift(); + activityCount -= 1; + if (removed?.status === 'error') { + failedActivityCount -= 1; + } else if (removed?.status === 'partial') { + partialActivityCount -= 1; + } } if (overflowActivityStartIndex == null || activity.startIndex > overflowActivityStartIndex) { overflowActivityStartIndex = activity.startIndex; @@ -828,6 +832,7 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity childLabelIndex, toolCallIds, startIndex: _startIndex, + partitionStartIndex: _partitionStartIndex, unresolvedToolStartIndex: _unresolvedToolStartIndex, ...activity }) => { @@ -882,7 +887,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity * answer content. */ const resolvedActivities = activities.map((activity) => { const toolStart = findTrackedToolStart(currentParts, activity); - return toolStart != null ? { ...activity, startIndex: toolStart } : activity; + return toolStart != null + ? { ...activity, startIndex: Math.max(toolStart, activity.partitionStartIndex ?? 0) } + : activity; }); const closesBeforeBoundary = (activity: TrackedActivity): boolean => { if (requestedEndIndex == null) { @@ -895,12 +902,30 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity return false; } const materializedStart = findMaterializedActivityStart(currentParts, activity); - return materializedStart == null || materializedStart < requestedEndIndex; + const materializedToolIndices = findTrackedToolIndices(currentParts, activity); + return materializedToolIndices.length > 0 + ? materializedToolIndices.every((index) => index < requestedEndIndex) + : materializedStart == null || materializedStart < requestedEndIndex; }; const snapshot = resolvedActivities.filter(closesBeforeBoundary); - const retainedActivities = resolvedActivities.filter( - (activity) => !closesBeforeBoundary(activity), - ); + const reanchorRetainedActivity = (activity: TrackedActivity): TrackedActivity => { + if (requestedEndIndex == null) { + return activity; + } + const firstRetainedToolIndex = findTrackedToolIndices(currentParts, activity).find( + (index) => index >= requestedEndIndex, + ); + return firstRetainedToolIndex != null + ? { + ...activity, + startIndex: firstRetainedToolIndex, + partitionStartIndex: firstRetainedToolIndex, + } + : activity; + }; + const retainedActivities = resolvedActivities + .filter((activity) => !closesBeforeBoundary(activity)) + .map(reanchorRetainedActivity); const overflowCount = Math.max(0, activityCount - resolvedActivities.length); const overflowFailedCount = Math.max( 0, @@ -922,7 +947,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity ? resolvedOverflowActivities.filter(closesBeforeBoundary) : []; const retainedOverflowActivities = hasExactOverflowActivities - ? resolvedOverflowActivities.filter((activity) => !closesBeforeBoundary(activity)) + ? resolvedOverflowActivities + .filter((activity) => !closesBeforeBoundary(activity)) + .map(reanchorRetainedActivity) : []; const overflowCloses = hasExactOverflowActivities ? closingOverflowActivities.length > 0 @@ -1176,10 +1203,17 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity } } } + const batchToolIndices = definedPartIndices(parts).filter((index) => { + const part = parts[index]; + return ( + part?.type === ContentTypes.TOOL_CALL && + typeof part.tool_call?.id === 'string' && + ids.has(part.tool_call.id) + ); + }); if ( - batchStartIndex != null && - emittedContentRanges.some( - ({ start, end }) => batchStartIndex >= start && batchStartIndex < end, + batchToolIndices.some((toolIndex) => + emittedContentRanges.some(({ start, end }) => toolIndex >= start && toolIndex < end), ) ) { pendingReasoning.delete(reasoningKey);