diff --git a/apps/website/content/docs/chat/api/api-docs.json b/apps/website/content/docs/chat/api/api-docs.json index 29f8a8a73..bdc922800 100644 --- a/apps/website/content/docs/chat/api/api-docs.json +++ b/apps/website/content/docs/chat/api/api-docs.json @@ -8783,7 +8783,7 @@ { "name": "completeDelivery", "kind": "function", - "description": "", + "description": "Creates a terminal delivery state for an existing response-attempt generation.", "signature": "completeDelivery(generation: string, outcome: TOutcome): object", "params": [ { @@ -9310,7 +9310,7 @@ { "name": "markdownDocument", "kind": "function", - "description": "", + "description": "Creates the atomic document consumed by the streaming Markdown renderer.", "signature": "markdownDocument(content: string, delivery: MessageDelivery, suffix: string): StreamingMarkdownDocument", "params": [ { @@ -9572,7 +9572,7 @@ { "name": "staticDelivery", "kind": "function", - "description": "", + "description": "Creates a successful terminal delivery state for an already-complete message.", "signature": "staticDelivery(messageId: string): object", "params": [ { @@ -9610,7 +9610,7 @@ { "name": "streamingDelivery", "kind": "function", - "description": "", + "description": "Creates the active delivery state for one response-attempt generation.", "signature": "streamingDelivery(generation: string): object", "params": [ { diff --git a/libs/ag-ui/src/lib/reducer.spec.ts b/libs/ag-ui/src/lib/reducer.spec.ts index 46988114f..13445db97 100644 --- a/libs/ag-ui/src/lib/reducer.spec.ts +++ b/libs/ag-ui/src/lib/reducer.spec.ts @@ -18,6 +18,7 @@ interface TestDeliveryRun { generation: string; baselineMessageIds: Set; ownedMessageIds: Set; + snapshotReplacementIds: Set; currentAssistantMessageId?: string; eligibleBaselineAssistantId?: string; protocolRunId?: string; @@ -46,6 +47,7 @@ function makeStore(generation = 'run-generation-1'): TestStore { generation, baselineMessageIds: new Set(), ownedMessageIds: new Set(), + snapshotReplacementIds: new Set(), }, allocateDeliveryGeneration: (scope: string) => `${generation}:${scope}:${++activitySequence}`, } as TestStore; @@ -150,7 +152,7 @@ describe('reduceEvent', () => { expect(store.deliveryRun?.currentAssistantMessageId).toBe('m2'); }); - it('does not mutate completed prior assistant content from a stale snapshot', () => { + it('publishes an authoritative rewrite of a completed step as a new completed generation', () => { const store = makeStore(); reduceEvent({ type: 'TEXT_MESSAGE_START', messageId: 'm1' } as any, store); reduceEvent({ type: 'TEXT_MESSAGE_CONTENT', messageId: 'm1', delta: 'first step' } as any, store); @@ -159,19 +161,52 @@ describe('reduceEvent', () => { reduceEvent({ type: 'MESSAGES_SNAPSHOT', messages: [ - { id: 'm1', role: 'assistant', content: 'stale rewrite' }, + { id: 'm1', role: 'assistant', content: 'authoritative rewrite' }, { id: 'm2', role: 'assistant', content: 'current snapshot' }, ], } as any, store); + const replacement = store.messages().find(message => message.id === 'm1'); + expect(replacement).toMatchObject({ + content: 'authoritative rewrite', + delivery: { phase: 'complete', outcome: 'success' }, + }); + const replacementGeneration = replacement?.delivery.generation; + expect(replacementGeneration).not.toBe('run-generation-1'); + expect(store.messages().find(message => message.id === 'm2')).toMatchObject({ + content: 'current snapshot', + delivery: streamingDelivery('run-generation-1'), + }); + + reduceEvent({ type: 'TEXT_MESSAGE_START', messageId: 'm1' } as any, store); + reduceEvent({ type: 'TEXT_MESSAGE_CONTENT', messageId: 'm1', delta: ' stale' } as any, store); expect(store.messages().find(message => message.id === 'm1')).toMatchObject({ - content: 'first step', - delivery: completeDelivery('run-generation-1', 'success'), + content: 'authoritative rewrite', + delivery: { generation: replacementGeneration, phase: 'complete', outcome: 'success' }, + }); + + const rewrittenSnapshot = { + type: 'MESSAGES_SNAPSHOT', + messages: [ + { id: 'm1', role: 'assistant', content: 'second authoritative rewrite' }, + ], + } as any; + reduceEvent(rewrittenSnapshot, store); + + const secondReplacement = store.messages().find(message => message.id === 'm1'); + expect(secondReplacement).toMatchObject({ + content: 'second authoritative rewrite', + delivery: { phase: 'complete', outcome: 'success' }, }); + expect(secondReplacement?.delivery.generation).not.toBe(replacementGeneration); expect(store.messages().find(message => message.id === 'm2')).toMatchObject({ content: 'current snapshot', delivery: streamingDelivery('run-generation-1'), }); + + reduceEvent(rewrittenSnapshot, store); + expect(store.messages().find(message => message.id === 'm1')?.delivery.generation) + .toBe(secondReplacement?.delivery.generation); }); it('preserves an omitted current assistant so later deltas still apply', () => { diff --git a/libs/ag-ui/src/lib/reducer.ts b/libs/ag-ui/src/lib/reducer.ts index a6c3f2bd8..d2aa160b8 100644 --- a/libs/ag-ui/src/lib/reducer.ts +++ b/libs/ag-ui/src/lib/reducer.ts @@ -65,6 +65,7 @@ export interface ReducerDeliveryRun { generation: string; baselineMessageIds: Set; ownedMessageIds: Set; + snapshotReplacementIds: Set; eligibleBaselineAssistantId?: string; currentAssistantMessageId?: string; protocolRunId?: string; @@ -339,10 +340,13 @@ export function reduceEvent(event: BaseEvent, store: ReducerStore): void { const snapshotToolCalls: ToolCall[] = []; const messages: Message[] = raw.map((m) => { const previous = previousById.get(m.id); - const completedOwnedMessage = run - && run.ownedMessageIds.has(m.id) - && previous?.delivery.generation === run.generation - && previous.delivery.phase === 'complete' + const completedMessage = m.id !== activeCanonicalAssistantId + && previous?.delivery.phase === 'complete' + && ( + !run + || run.ownedMessageIds.has(m.id) + || run.snapshotReplacementIds.has(m.id) + ) ? previous : undefined; let delivery = previous?.delivery ?? staticDelivery(m.id); @@ -363,7 +367,21 @@ export function reduceEvent(event: BaseEvent, store: ReducerStore): void { const { toolCalls: _dropped, ...rest } = m; snapshotMessage = { ...rest, toolCallIds: ids } as unknown as Omit; } - if (completedOwnedMessage) return completedOwnedMessage; + if (completedMessage) { + const snapshotChanged = completedMessage.content !== snapshotMessage.content + || !sameStringArray(completedMessage.toolCallIds, snapshotMessage.toolCallIds); + if (!snapshotChanged) return completedMessage; + + run?.snapshotReplacementIds.add(m.id); + return { + ...completedMessage, + ...snapshotMessage, + delivery: completeDelivery( + store.allocateDeliveryGeneration(`snapshot:${m.id}`), + 'success', + ), + } as Message; + } if (run && (run.ownedMessageIds.has(m.id) || m.id === canonicalAssistantId)) { delivery = delivery.generation === run.generation && delivery.phase === 'complete' ? delivery @@ -518,6 +536,12 @@ function ownAssistantMessage(store: ReducerStore, id: string) { return streamingDelivery(run.generation); } +function sameStringArray(left?: string[], right?: string[]): boolean { + if (left === right) return true; + if (!left || !right || left.length !== right.length) return false; + return left.every((value, index) => value === right[index]); +} + function resolveCanonicalAssistantId( messages: readonly AgUiSnapshotMessage[], run: ReducerDeliveryRun | null, diff --git a/libs/ag-ui/src/lib/to-agent.conformance.spec.ts b/libs/ag-ui/src/lib/to-agent.conformance.spec.ts index c04bd8cdd..dc73c03ce 100644 --- a/libs/ag-ui/src/lib/to-agent.conformance.spec.ts +++ b/libs/ag-ui/src/lib/to-agent.conformance.spec.ts @@ -84,6 +84,7 @@ describe('AG-UI reducer — reasoning-fixture conformance', () => { generation: 'reasoning-fixture-run', baselineMessageIds: new Set(), ownedMessageIds: new Set(), + snapshotReplacementIds: new Set(), }, allocateDeliveryGeneration: (scope: string) => `reasoning-fixture:${scope}`, }; diff --git a/libs/ag-ui/src/lib/to-agent.ts b/libs/ag-ui/src/lib/to-agent.ts index 5a729f5d6..d8c19f2b1 100644 --- a/libs/ag-ui/src/lib/to-agent.ts +++ b/libs/ag-ui/src/lib/to-agent.ts @@ -185,6 +185,7 @@ export function toAgent(source: AbstractAgent, options: ToAgentOptions = {}): Ag generation: allocateDeliveryGeneration('run'), baselineMessageIds: new Set(store.messages().map(message => message.id)), ownedMessageIds: new Set(), + snapshotReplacementIds: new Set(), eligibleBaselineAssistantId: allowBaselineTail ? getTailAssistantMessageId(store.messages()) : undefined, diff --git a/libs/chat/src/lib/agent/message-delivery.ts b/libs/chat/src/lib/agent/message-delivery.ts index 6f6194f9e..02938b1f8 100644 --- a/libs/chat/src/lib/agent/message-delivery.ts +++ b/libs/chat/src/lib/agent/message-delivery.ts @@ -22,10 +22,12 @@ export type MessageDelivery = readonly outcome: CompleteOutcome; }; +/** Creates the active delivery state for one response-attempt generation. */ export function streamingDelivery(generation: string) { return { generation, phase: 'streaming' } as const satisfies MessageDelivery; } +/** Creates a terminal delivery state for an existing response-attempt generation. */ export function completeDelivery( generation: string, outcome: TOutcome, @@ -34,6 +36,7 @@ export function completeDelivery( return delivery satisfies MessageDelivery; } +/** Creates a successful terminal delivery state for an already-complete message. */ export function staticDelivery(messageId: string) { return completeDelivery(messageId, 'success'); } diff --git a/libs/chat/src/lib/streaming/streaming-markdown.component.ts b/libs/chat/src/lib/streaming/streaming-markdown.component.ts index 38f83891c..531467701 100644 --- a/libs/chat/src/lib/streaming/streaming-markdown.component.ts +++ b/libs/chat/src/lib/streaming/streaming-markdown.component.ts @@ -31,6 +31,7 @@ export interface StreamingMarkdownDocument { readonly content: string; } +/** Creates the atomic document consumed by the streaming Markdown renderer. */ export function markdownDocument( content: string, delivery: MessageDelivery,