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

/** Web presentation fold joining transient Assistant frames to one durable v2 settlement. */

import type {
  SessionAssistantStreamBaseline,
  SessionAssistantStreamFrame,
} from '../../types.ts'
import { expandAssistantStream } from '@deepseek-ai/dsh-llm/assistant-stream'
import type { AssistantStreamRecord } from '@deepseek-ai/dsh-llm/assistant-stream'
import type {
  SessionEventLikeEntry,
  SessionLiveEventEntry,
  SessionTransientEventEntry,
} from '../contract/events.ts'

interface ActiveAttempt {
  readonly attemptId: string
  readonly startedAfterSeq: number
  readonly turn: number
  readonly step: number
  nextIndex: number
}

/** One Web publication decision from the assistant stream fold. */
export type ClientAssistantStreamResult =
  | { readonly type: 'publish'; readonly entry: SessionLiveEventEntry }
  | {
    readonly type: 'settlement'
    readonly attemptId: string
    readonly entry: SessionLiveEventEntry
  }
  | { readonly type: 'transient'; readonly entry: SessionTransientEventEntry }
  | { readonly type: 'rebaseline' }
  | undefined

/** Keeps transient Assistant presentation behind one settlement-aware interface. */
export class ClientAssistantStream {
  private activeAttempt: ActiveAttempt | undefined
  private readonly pending = new Map<number, SessionLiveEventEntry>()
  private publishedSeqs = new Set<number>()
  private durableCursor = -1
  private transientInGap = 0

  /**
   * Replace the durable Web window and adopt an optional reconnect baseline.
   * @param entries - durable entries in the replacement window.
   * @param baseline - compact prefix for an Assistant attempt that is still live.
   * @returns immediately visible durable entries plus reconstructed transient chunks.
   */
  replace(
    entries: readonly SessionEventLikeEntry[],
    baseline?: SessionAssistantStreamBaseline,
  ): readonly SessionEventLikeEntry[] {
    this.pending.clear()
    this.transientInGap = 0
    this.activeAttempt = undefined
    const opening = baseline?.activeAttempt
    if (opening !== undefined) {
      this.activeAttempt = {
        attemptId: String(opening.attemptId),
        startedAfterSeq: opening.startedAfterSeq,
        turn: opening.turn,
        step: opening.step,
        nextIndex: opening.nextIndex,
      }
    }
    const visible: SessionEventLikeEntry[] = [...entries]
    this.publishedSeqs = new Set(visible.map(entry => entry.event.seq))
    this.durableCursor = visible.reduce((cursor, entry) => Math.max(cursor, entry.event.seq), -1)
    if (opening !== undefined) {
      for (const [index, member] of expandAssistantStream(
        opening.stream as unknown as readonly AssistantStreamRecord[],
      ).entries()) {
        this.transientInGap += 1
        visible.push({
          type: 'transient',
          event: {
            type: 'assistant/live-chunk',
            seq: this.durableCursor + 1 - 1 / (this.transientInGap + 1),
            time: member.time,
            data: {
              attemptId: opening.attemptId,
              turn: opening.turn,
              step: opening.step,
              chunk: member.chunk,
            },
          },
        })
        if (index + 1 >= opening.nextIndex) break
      }
    }
    return visible
  }

  /**
   * Stage one durable v2 settlement while its matching live attempt is open.
   * @param entry - newly followed durable entry.
   * @returns a publication decision, or `undefined` when no entry becomes visible.
   */
  acceptDurable(entry: SessionLiveEventEntry): ClientAssistantStreamResult {
    const event = entry.event
    this.durableCursor = Math.max(this.durableCursor, event.seq)
    this.transientInGap = 0
    if (this.attemptForSettlement(event) !== undefined) {
      this.pending.set(event.seq, entry)
      return undefined
    }
    return this.publish(entry)
  }

  /**
   * Fold one dense transient frame and release its named durable settlement.
   * @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':
        this.pending.clear()
        this.activeAttempt = {
          attemptId: String(frame.attemptId),
          startedAfterSeq: frame.startedAfterSeq,
          turn: frame.turn,
          step: frame.step,
          nextIndex: 0,
        }
        return undefined
      case 'chunk': {
        const attempt = this.activeAttempt
        if (attempt === undefined
          || attempt.attemptId !== String(frame.attemptId)
          || frame.index !== attempt.nextIndex) return { type: 'rebaseline' }
        attempt.nextIndex += 1
        this.transientInGap += 1
        return {
          type: 'transient',
          entry: {
            type: 'transient',
            event: {
              type: 'assistant/live-chunk',
              seq: this.durableCursor + 1 - 1 / (this.transientInGap + 1),
              time: frame.time,
              data: {
                attemptId: frame.attemptId,
                turn: attempt.turn,
                step: attempt.step,
                chunk: frame.chunk as never,
              },
            },
          },
        }
      }
      case 'end': {
        const attempt = this.activeAttempt
        this.activeAttempt = undefined
        if (attempt === undefined || attempt.attemptId !== String(frame.attemptId)) {
          return { type: 'rebaseline' }
        }
        if (frame.index !== attempt.nextIndex) return { type: 'rebaseline' }
        if (frame.outcome.kind === 'abandoned') {
          return this.pending.size === 0 ? undefined : { type: 'rebaseline' }
        }
        if (this.publishedSeqs.has(frame.outcome.seq)) return undefined
        const entry = this.pending.get(frame.outcome.seq)
        if (entry === undefined
          || entry.event.type !== frame.outcome.eventType
          || entry.event.data.turn !== attempt.turn
          || entry.event.data.step !== attempt.step) {
          return { type: 'rebaseline' }
        }
        this.pending.delete(frame.outcome.seq)
        this.publishedSeqs.add(entry.event.seq)
        return { type: 'settlement', attemptId: attempt.attemptId, entry }
      }
    }
  }

  private attemptForSettlement(
    event: SessionLiveEventEntry['event'],
  ): ActiveAttempt | undefined {
    const attempt = this.activeAttempt
    if (attempt === undefined
      || (event.type !== 'assistant/message' && event.type !== 'assistant/attempt')
      || (event.type === 'assistant/message' && event.surfaceOp !== 'append')
      || event.seq <= attempt.startedAfterSeq
      || attempt.turn !== event.data.turn
      || attempt.step !== event.data.step) return undefined
    return attempt
  }

  private publish(entry: SessionLiveEventEntry): ClientAssistantStreamResult {
    this.publishedSeqs.add(entry.event.seq)
    return { type: 'publish', entry }
  }
}