fix(chat): retain preparing tools and ordering until step completion

packages/api/session-controller/src/client/sessions/assistant-stream.ts

@@ -24,7 +24,12 @@ interface ActiveAttempt {
 
 /** One Web publication decision from the assistant stream fold. */
 export type ClientAssistantStreamResult =
-  | { readonly type: 'publish'; readonly entry: SessionLiveEventEntry }
+  | {
+    readonly type: 'publish'
+    readonly entry: SessionLiveEventEntry
+    /** Retire the successful attempt's transient rows after publishing its Step end. */
+    readonly retireAttemptId?: LlmAttemptId
+  }
   | {
     readonly type: 'settlement'
     readonly attemptId: LlmAttemptId
@@ -38,6 +43,7 @@ export type ClientAssistantStreamResult =
 /** Keeps transient Assistant presentation behind one settlement-aware interface. */
 export class ClientAssistantStream {
   private activeAttempt: ActiveAttempt | undefined
+  private retainedAttempt: Pick<ActiveAttempt, 'attemptId' | 'turn' | 'step'> | undefined
   private readonly pending = new Map<number, SessionAssistantSettlementEntry>()
   private publishedSeqs = new Set<number>()
   private durableCursor = -1
@@ -56,6 +62,7 @@ export class ClientAssistantStream {
     this.pending.clear()
     this.transientInGap = 0
     this.activeAttempt = undefined
+    this.retainedAttempt = undefined
     const opening = baseline?.activeAttempt
     if (opening !== undefined) {
       this.activeAttempt = {
@@ -114,13 +121,16 @@ export class ClientAssistantStream {
 
   /**
    * Fold one dense transient frame and release its named durable settlement.
+   * Successful messages retain their transient rows until the owning Step ends;
+   * interrupted messages, failed attempts, and abandonment retire them immediately.
    * @param frame - next Assistant stream frame received by the follow connection.
    * @returns a transient, publication, or rebaseline decision, or `undefined` when no entry becomes visible.
    */
   acceptFrame(frame: SessionAssistantStreamFrame): ClientAssistantStreamResult {
     switch (frame.type) {
       case 'start':
-        if (this.activeAttempt !== undefined || this.pending.size > 0) return { type: 'rebaseline' }
+        if (this.activeAttempt !== undefined || this.retainedAttempt !== undefined
+          || this.pending.size > 0) return { type: 'rebaseline' }
         this.pending.clear()
         this.activeAttempt = {
           attemptId: frame.attemptId,
@@ -176,6 +186,10 @@ export class ClientAssistantStream {
           return { type: 'rebaseline' }
         }
         this.pending.delete(frame.outcome.seq)
+        if (entry.event.type === 'assistant/message' && entry.event.data.interrupted !== true) {
+          this.retainedAttempt = { attemptId: attempt.attemptId, turn: attempt.turn, step: attempt.step }
+          return this.publish(entry)
+        }
         this.publishedSeqs.add(entry.event.seq)
         return { type: 'settlement', attemptId: attempt.attemptId, entry }
       }
@@ -196,6 +210,12 @@ export class ClientAssistantStream {
 
   private publish(entry: SessionLiveEventEntry): ClientAssistantStreamResult {
     this.publishedSeqs.add(entry.event.seq)
+    const retained = this.retainedAttempt
+    if (retained !== undefined && entry.event.type === 'step/end'
+      && entry.event.data.turn === retained.turn && entry.event.data.step === retained.step) {
+      this.retainedAttempt = undefined
+      return { type: 'publish', entry, retireAttemptId: retained.attemptId }
+    }
     return { type: 'publish', entry }
   }
 }

packages/api/session-controller/src/client/sessions/session.ts

@@ -723,8 +723,10 @@ export class Session implements SessionFace {
       this.notifier.markDirty()
       return
     }
-    if (result?.type === 'publish' && this.appendLive(result.entry)) {
-      this.notifier.markDirty()
+    if (result?.type === 'publish') {
+      const changed = this.appendLive(result.entry)
+      if (result.retireAttemptId !== undefined) this.eventSource.settleAssistant(result.retireAttemptId)
+      if (changed || result.retireAttemptId !== undefined) this.notifier.markDirty()
     } else if (result?.type === 'transient') {
       this.eventSource.append(result.entry)
       this.notifier.markDirty()

packages/api/session-controller/tests/assistant-stream.client.spec.ts

@@ -136,14 +136,48 @@ describe('ClientAssistantStream', () => {
     }
   })
 
-  it('stages one owned settlement and releases it from the matching end frame', () => {
+  it('publishes a successful message at stream end and retires its chunks at its Step end', () => {
     const stream = opened()
     const durable = messageEvent(2)
     expect(stream.acceptDurable(durable)).toBeUndefined()
     expect(stream.acceptFrame(chunkFrame(0))).toEqual(expect.objectContaining({ type: 'transient' }))
     expect(stream.acceptFrame(end(1, {
       kind: 'committed', eventType: 'assistant/message', seq: 2,
-    }))).toEqual({ type: 'settlement', attemptId: String(ATTEMPT), entry: durable })
+    }))).toEqual({ type: 'publish', entry: durable })
+    for (const [turn, step] of [[2, 1], [1, 2]]) {
+      const unrelated = entry({ type: 'step/end', seq: SessionSeq(3), time: 3, data: { turn: turn!, step: step! } })
+      expect(stream.acceptDurable(unrelated)).toEqual({ type: 'publish', entry: unrelated })
+    }
+    const closed = entry({ type: 'step/end', seq: SessionSeq(4), time: 4, data: { turn: 1, step: 1 } })
+    expect(stream.acceptDurable(closed)).toEqual({ type: 'publish', entry: closed, retireAttemptId: ATTEMPT })
+    expect(stream.acceptDurable(ordinary(5))).toEqual({ type: 'publish', entry: ordinary(5) })
+    expect(stream.acceptFrame(start(LlmAttemptId('session:2'), 4))).toBeUndefined()
+  })
+
+  it('clears retained chunks on replacement and refuses a new attempt before Step completion', () => {
+    const stream = opened()
+    const durable = messageEvent(2)
+    stream.acceptDurable(durable)
+    stream.acceptFrame(end(0, { kind: 'committed', eventType: 'assistant/message', seq: 2 }))
+    expect(stream.acceptFrame(start(LlmAttemptId('session:2')))).toEqual({ type: 'rebaseline' })
+    expect(stream.replace([durable])).toEqual([durable])
+    const closed = entry({ type: 'step/end', seq: SessionSeq(3), time: 3, data: { turn: 1, step: 1 } })
+    expect(stream.acceptDurable(closed)).toEqual({ type: 'publish', entry: closed })
+    expect(stream.acceptFrame(start(LlmAttemptId('session:2'), 3))).toBeUndefined()
+  })
+
+  it('retires failed attempts and interrupted messages at stream end', () => {
+    const message = messageEvent(2)
+    if (message.event.type !== 'assistant/message') throw new Error('expected assistant message')
+    const interrupted = entry({ ...message.event, data: { ...message.event.data, interrupted: true as const } })
+    for (const durable of [attemptEvent(2), interrupted]) {
+      const stream = opened()
+      stream.acceptDurable(durable)
+      expect(stream.acceptFrame(end(0, {
+        kind: 'committed', eventType: durable.event.type as 'assistant/message' | 'assistant/attempt', seq: 2,
+      }))).toEqual({ type: 'settlement', attemptId: ATTEMPT, entry: durable })
+      expect(stream.acceptFrame(start(LlmAttemptId('session:2'), 2))).toBeUndefined()
+    }
   })
 
   it('rebaselines duplicate durable settlements or starts', () => {

packages/api/session-controller/tests/sessions-service.client.spec.ts

@@ -164,7 +164,7 @@ describe('scope tree', () => {
     expect(b.mock.remote.session.projections).not.toHaveBeenCalled()
   })
 
-  it('publishes transient Assistant chunks and the named durable v2 settlement through one event source', async ({ bench }) => {
+  it('publishes the final Assistant message before retiring transient chunks at Step end', async ({ bench }) => {
     const b = bench()
     await feedList(b, [{ id: 's1' }])
     using _reference = b.svc.retain(sid('s1'), { source: 'controllerOperation' })
@@ -232,13 +232,20 @@ describe('scope tree', () => {
     })
     await b.mock.streams.drained(FOLLOW)
     await vi.waitFor(() => {
-      expect(binding.eventSource.getSnapshot().entries).toHaveLength(1)
+      expect(binding.eventSource.getSnapshot().entries).toHaveLength(2)
     })
 
     expect(publications).toEqual([
       ['assistant/live-chunk'],
-      ['assistant/message'],
+      ['assistant/live-chunk', 'assistant/message'],
     ])
+    b.mock.streams.push(FOLLOW, {
+      type: 'event', event: { type: 'step/end', seq: 1, time: 3, data: { turn: 1, step: 1 } },
+    })
+    await b.mock.streams.drained(FOLLOW)
+    expect(binding.eventSource.getSnapshot().entries.map(entry => entry.event.type))
+      .toEqual(['assistant/message', 'step/end'])
+    expect(binding.eventSource.getSnapshot().change).toEqual({ kind: 'settle-assistant', attemptId })
     dispose()
   })
 
@@ -373,10 +380,10 @@ describe('scope tree', () => {
     await b.mock.streams.drained(FOLLOW)
     await vi.waitFor(() => {
       expect(binding.eventSource.getSnapshot().entries.map(entry => entry.event.type))
-        .toEqual(['assistant/message', 'assistant/message'])
+        .toEqual(['assistant/message', 'assistant/live-chunk', 'assistant/message'])
     })
     expect(binding.eventSource.getSnapshot().change).toEqual({
-      kind: 'settle-assistant', attemptId: String(attemptId), entry: currentMessage,
+      kind: 'append', entries: [currentMessage],
     })
   })
 

packages/client/ui-chat/src/client/conversation-nodes/assistant.ts

@@ -259,7 +259,8 @@ function projectAssistant(context: ConversationNodeContext<AssistantState>): Ass
   const status = settled?.interrupted === true
     ? 'interrupted'
     : settled === undefined ? 'running' : 'settled'
-  const anchorSeq = settled?.seq ?? state.firstVisibleSeq ?? context.matches[0]?.event.seq ?? 0
+  const anchorSeq = (settled?.interrupted === true ? settled.seq : state.firstVisibleSeq ?? settled?.seq)
+    ?? context.matches[0]?.event.seq ?? 0
   const time = settled?.time ?? state.firstVisibleTime ?? context.matches[0]?.event.time ?? 0
   return {
     anchorSeq,
@@ -341,7 +342,9 @@ export const assistantDefinition: ConversationNodeDefinition<AssistantState> = {
     const settled = data.finalNode
     const visible = settled === undefined ? state.visibleBlocks > 0 : hasVisibleContent(data.blocks)
     if (settled === undefined && !visible && current == null) return null
-    const anchorSeq = settled?.seq ?? state.firstVisibleSeq ?? context.matches[0]?.event.seq ?? 0
+    // A successful message retains its live anchor alongside pending Tool calls.
+    const anchorSeq = (settled?.interrupted === true ? settled.seq : state.firstVisibleSeq ?? settled?.seq)
+      ?? context.matches[0]?.event.seq ?? 0
     return chatNode(context, 'assistant-step', anchorSeq, data, {
       visibility: settled?.interrupted === true || visible ? 'visible' : 'hidden',
     })

packages/client/ui-chat/tests/tool-preparation.client.spec.ts

@@ -25,22 +25,22 @@ const opening = [
   entry({ type: 'step/start', seq: SessionSeq(2), time: 2, data: { turn: 1, step: 1 } }),
 ]
 
-function delta(seq: number, id = first, argumentsDelta = '', name: string | null = 'write'): SessionEventLikeEntry {
+function delta(seq: number, id = first, argumentsDelta = '', name: string | null = 'write', index = id === first ? 0 : 1): SessionEventLikeEntry {
   return {
     type: 'transient',
     event: {
       type: 'assistant/live-chunk', seq, time: seq,
       data: { turn: 1, step: 1, attemptId, chunk: {
-        type: 'tool-call-delta', index: id === first ? 0 : 1, id, argumentsDelta,
+        type: 'tool-call-delta', index, id, argumentsDelta,
         ...(name === null ? {} : { name }),
       } },
     },
   }
 }
 
-function call(seq: number, callId = first): SessionLiveEventEntry {
+function call(seq: number, callId = first, name = 'write', argumentsRaw = args): SessionLiveEventEntry {
   return entry({ type: 'tool/call', seq: SessionSeq(seq), time: seq,
-    data: { turn: 1, step: 1, callId, name: 'write', arguments: args } })
+    data: { turn: 1, step: 1, callId, name, arguments: argumentsRaw } })
 }
 
 function result(seq: number, callId = first): SessionLiveEventEntry {
@@ -49,18 +49,20 @@ function result(seq: number, callId = first): SessionLiveEventEntry {
   } })
 }
 
-function settlement(): SessionAssistantSettlementEntry {
+function settlement(ids = [first, second], name = 'write', argumentsRaw = args, text = ''): SessionAssistantSettlementEntry {
   const stream = new AssistantStreamAccumulator()
-  for (const [index, id] of [first, second].entries()) {
-    const chunk: StreamChunk = { type: 'tool-call-delta', index, id, name: 'write', argumentsDelta: args }
-    stream.push({ time: 3 + index, chunk })
+  if (text !== '') stream.push({ time: 3, chunk: { type: 'text-delta', index: 0, text } })
+  for (const [index, id] of ids.entries()) {
+    const chunk: StreamChunk = { type: 'tool-call-delta', index: index + Number(text !== ''), id, name, argumentsDelta: argumentsRaw }
+    stream.push({ time: 3, chunk })
   }
   return { type: 'event', event: {
     type: 'assistant/message', seq: SessionSeq(5), time: 5, surfaceOp: 'append', data: {
       turn: 1, step: 1, stream: [...stream.snapshot()],
-      message: createAssistantMessage({ source: { provider: 'test', model: 'test' }, content: [first, second].map(id => ({
-        type: 'tool-call', id, name: 'write', arguments: args,
-      })) }),
+      message: createAssistantMessage({ source: { provider: 'test', model: 'test' }, content: [
+        ...text === '' ? [] : [{ type: 'text' as const, text }],
+        ...ids.map(id => ({ type: 'tool-call' as const, id, name, arguments: argumentsRaw })),
+      ] }),
     },
   } }
 }
@@ -130,19 +132,62 @@ describe('Tool preparation and durable replay', () => {
     expect(h.phases()).toEqual([[first, 'result'], [second, 'result']])
   })
 
-  it('withdraws transient preparation on settlement and converges with a durable-only replacement', () => {
+  it('keeps five command preparations below the Assistant message through dispatch and Step end', () => {
+    const h = harness()
+    const ids = Array.from({ length: 5 }, (_, index) => ToolCallId(`command-${index}`))
+    const command = '{"command":"echo hello"}'
+    const text = 'Running five commands'
+    h.assembler.append({ type: 'transient', event: {
+      type: 'assistant/live-chunk', seq: 2.01, time: 3,
+      data: { turn: 1, step: 1, attemptId, chunk: { type: 'text-delta', index: 0, text } },
+    } })
+    for (const [index, id] of ids.entries()) h.assembler.append(delta(2.1 + index / 10, id, command, 'bash', index + 1))
+    const keys = h.tools().map(node => node.key)
+    const groups = h.assembler.grouped('chat')!
+    const group = groups.entries.find(entry => entry.kind === 'group')!
+    const order = groups.entries
+    expect(order[0]).toMatchObject({ kind: 'node', groupPart: 'response' })
+    expect(order[1]).toEqual(group)
+    const message = settlement(ids, 'bash', command, text)
+    const history = [...opening, message]
+    h.assembler.append(message)
+    expect(h.phases()).toEqual(ids.map(id => [id, 'preparing']))
+    expect(groups.entries).toEqual(order)
+    for (const [index, id] of ids.entries()) {
+      const started = call(6 + index * 2, id, 'bash', command)
+      h.assembler.append(started)
+      expect(h.phases()).toEqual(ids.map((callId, at) => [callId, at < index ? 'result' : at === index ? 'start' : 'preparing']))
+      expect(groups.entries).toEqual(order)
+      const completed = result(7 + index * 2, id)
+      h.assembler.append(completed)
+      history.push(started, completed)
+      expect(h.tools().map(node => node.key)).toEqual(keys)
+      expect(groups.entries).toEqual(order)
+    }
+    const closed = entry({ type: 'step/end', seq: SessionSeq(16), time: 16, data: { turn: 1, step: 1 } })
+    history.push(closed)
+    h.assembler.append(closed)
+    h.assembler.settleAssistant(attemptId)
+    const replay = harness(history)
+    expect(h.phases()).toEqual(ids.map(id => [id, 'result']))
+    expect(h.material()).toEqual(replay.material())
+    expect(groups.entries).toEqual(order)
+  })
+
+  it('hides undispatched preparations when a settled Step ends before every call starts', () => {
     const h = harness()
     h.assembler.append(delta(2.1))
-    const key = h.tools()[0]!.key
+    h.assembler.append(delta(2.2, second))
     const message = settlement()
-    h.assembler.settleAssistant(attemptId, message)
-    expect(h.phases()).toEqual([])
-    const history = [...opening, message, call(6), result(7), call(8, second), result(9, second)]
-    for (const event of history.slice(3)) h.assembler.append(event)
-    expect(h.tools()[0]?.key).toBe(key)
-    const replay = harness(history)
-    expect(h.phases()).toEqual([[first, 'result'], [second, 'result']])
-    expect(h.material()).toEqual(replay.material())
+    h.assembler.append(message)
+    h.assembler.append(call(6))
+    h.assembler.append(result(7))
+    expect(h.phases()).toEqual([[first, 'result'], [second, 'preparing']])
+    const closed = entry({ type: 'step/end', seq: SessionSeq(8), time: 8, data: { turn: 1, step: 1 } })
+    h.assembler.append(closed)
+    expect(h.phases()).toEqual([[first, 'result']])
+    h.assembler.settleAssistant(attemptId)
+    expect(h.material()).toEqual(harness([...opening, message, call(6), result(7), closed]).material())
   })
 
   it('does not create preparing rows from stored Assistant streams in history or older pages', () => {