@tanstack/ai-client 0.22.1 → 0.23.1

package/src/cleared-stream-tracker.ts

@@ -0,0 +1,151 @@
+import { getChunkRunId } from './connection-adapters'
+import type { StreamChunk } from '@tanstack/ai/client'
+import type { UIMessage } from './types'
+
+function getChunkToolCallId(chunk: StreamChunk): string | undefined {
+  return 'toolCallId' in chunk && typeof chunk.toolCallId === 'string'
+    ? chunk.toolCallId
+    : undefined
+}
+
+function getChunkMessageId(chunk: StreamChunk): string | undefined {
+  return 'messageId' in chunk && typeof chunk.messageId === 'string'
+    ? chunk.messageId
+    : undefined
+}
+
+function getChunkParentMessageId(chunk: StreamChunk): string | undefined {
+  return 'parentMessageId' in chunk && typeof chunk.parentMessageId === 'string'
+    ? chunk.parentMessageId
+    : undefined
+}
+
+/** Tracks stream chunks that must be ignored after the owning chat is cleared. */
+export class ClearedStreamTracker {
+  private readonly clearedMessageIds = new Set<string>()
+  private readonly clearedRunIds = new Set<string>()
+  private readonly ignoredActiveRunIds = new Set<string>()
+  private readonly clearedToolCallIds = new Set<string>()
+  private currentRunlessRunId: string | null = null
+
+  snapshotClear(context: {
+    messages: Array<UIMessage>
+    activeRunIds: Set<string>
+    currentRunId: string | null
+  }): void {
+    for (const message of context.messages) {
+      this.clearedMessageIds.add(message.id)
+    }
+    for (const runId of context.activeRunIds) {
+      this.clearedRunIds.add(runId)
+      this.ignoredActiveRunIds.add(runId)
+    }
+    if (context.currentRunId) {
+      this.clearedRunIds.add(context.currentRunId)
+      this.ignoredActiveRunIds.add(context.currentRunId)
+    }
+  }
+
+  shouldIgnoreChunk(chunk: StreamChunk): boolean {
+    const runId = getChunkRunId(chunk)
+    if (runId && this.clearedRunIds.has(runId)) {
+      if (chunk.type === 'RUN_STARTED') {
+        this.ignoredActiveRunIds.add(runId)
+        this.currentRunlessRunId = runId
+      }
+      this.markIgnoredChunkIds(chunk)
+      return true
+    }
+
+    if (runId && this.ignoredActiveRunIds.has(runId)) {
+      this.markIgnoredChunkIds(chunk)
+      return true
+    }
+
+    if (this.isRunlessChunkFromIgnoredRun(chunk)) {
+      this.markIgnoredChunkIds(chunk)
+      return true
+    }
+
+    const toolCallId = getChunkToolCallId(chunk)
+    if (toolCallId && this.clearedToolCallIds.has(toolCallId)) {
+      return true
+    }
+
+    const parentMessageId = getChunkParentMessageId(chunk)
+    if (parentMessageId && this.clearedMessageIds.has(parentMessageId)) {
+      if (toolCallId) {
+        this.clearedToolCallIds.add(toolCallId)
+      }
+      return true
+    }
+
+    const messageId = getChunkMessageId(chunk)
+    return Boolean(messageId && this.clearedMessageIds.has(messageId))
+  }
+
+  onRunStarted(runId: string): void {
+    this.currentRunlessRunId = runId
+  }
+
+  onRunSettled(runId: string): void {
+    this.ignoredActiveRunIds.delete(runId)
+    this.clearedRunIds.delete(runId)
+    if (this.currentRunlessRunId === runId) {
+      this.currentRunlessRunId =
+        this.ignoredActiveRunIds.values().next().value ?? null
+    }
+  }
+
+  onSessionRunError(): void {
+    this.ignoredActiveRunIds.clear()
+    this.currentRunlessRunId = null
+  }
+
+  resetActiveRuns(): void {
+    this.ignoredActiveRunIds.clear()
+    this.currentRunlessRunId = null
+  }
+
+  takeRunlessRunId(): string | null {
+    const runId = this.currentRunlessRunId
+    if (!runId) return null
+    this.ignoredActiveRunIds.delete(runId)
+    this.clearedRunIds.delete(runId)
+    this.currentRunlessRunId =
+      this.ignoredActiveRunIds.values().next().value ?? null
+    return runId
+  }
+
+  private markIgnoredChunkIds(chunk: StreamChunk): void {
+    const messageId = getChunkMessageId(chunk)
+    if (messageId) {
+      this.clearedMessageIds.add(messageId)
+    }
+    const toolCallId = getChunkToolCallId(chunk)
+    if (toolCallId) {
+      this.clearedToolCallIds.add(toolCallId)
+    }
+  }
+
+  private isRunlessChunkFromIgnoredRun(chunk: StreamChunk): boolean {
+    const runId = getChunkRunId(chunk)
+    if (runId || !this.currentRunlessRunId) return false
+    if (
+      !this.ignoredActiveRunIds.has(this.currentRunlessRunId) &&
+      !this.clearedRunIds.has(this.currentRunlessRunId)
+    ) {
+      return false
+    }
+    return (
+      chunk.type === 'TEXT_MESSAGE_START' ||
+      chunk.type === 'TEXT_MESSAGE_CONTENT' ||
+      chunk.type === 'TOOL_CALL_START' ||
+      chunk.type === 'TOOL_CALL_ARGS' ||
+      chunk.type === 'TOOL_CALL_END' ||
+      chunk.type === 'TOOL_CALL_RESULT' ||
+      chunk.type === 'MESSAGES_SNAPSHOT' ||
+      chunk.type === 'RUN_ERROR'
+    )
+  }
+}

package/src/client-persistor.ts

@@ -1,6 +1,20 @@
 import { getChunkRunId } from './connection-adapters'
 import type { StreamChunk } from '@tanstack/ai/client'
-import type { ChatClientPersistence, UIMessage } from './types'
+import type {
+  ChatClientPersistence,
+  ChatPersistedState,
+  ChatResumeSnapshot,
+  UIMessage,
+} from './types'
+
+/** Normalize a raw `getItem` result (legacy bare array or combined record). */
+function normalizePersistedState(
+  raw: ChatPersistedState | Array<UIMessage> | null | undefined,
+): ChatPersistedState | undefined {
+  if (Array.isArray(raw)) return { messages: raw }
+  if (raw && Array.isArray(raw.messages)) return raw
+  return undefined
+}
 
 // `StreamChunk` is a discriminated union; `toolCallId` / `messageId` /
 // `parentMessageId` exist on only some members. Narrow with `in` (matching
@@ -30,10 +44,12 @@ function getChunkParentMessageId(chunk: StreamChunk): string | undefined {
  *
  * Two responsibilities live here:
  *
- * 1. **Storage orchestration** — hydrate from `getItem(id)` on creation, save to
- *    `setItem(id, messages)` on every change through an ordered write queue, and
- *    `removeItem(id)` on clear. A generation counter discards stale writes when a
- *    removal or a newer conversation supersedes an in-flight async operation.
+ * 1. **Storage orchestration** — hydrate from `getItem(id)` on creation, save a
+ *    combined `{ messages, resume? }` record via `setItem` on every change
+ *    through an ordered write queue, and `removeItem(id)` on clear (or when
+ *    both the transcript and resume pointer are empty). A generation counter
+ *    discards stale writes when a removal or a newer conversation supersedes
+ *    an in-flight async operation.
  * 2. **Clear-during-stream suppression** — when a conversation is cleared while a
  *    stream is still producing, late chunks for the cleared run(s) must not
  *    repopulate the now-empty state. The persistor tracks the cleared ids and
@@ -51,6 +67,10 @@ export class ChatPersistor {
   // Bumped on every message change; lets an in-flight async hydration detect
   // that the message list moved on and avoid clobbering it.
   private messagesGeneration = 0
+  // Latest messages + resume snapshot, written together as one combined record
+  // so a full page reload restores both from a single adapter key.
+  private lastMessages: Array<UIMessage> = []
+  private lastResume: ChatResumeSnapshot | null = null
 
   // --- clear-during-stream suppression state ---
   private readonly clearedMessageIds = new Set<string>()
@@ -63,51 +83,91 @@ export class ChatPersistor {
     private readonly adapter: ChatClientPersistence,
     private readonly id: string,
     private readonly applyMessages: (messages: Array<UIMessage>) => void,
+    private readonly applyResume?: (snapshot: ChatResumeSnapshot) => void,
   ) {}
 
+  /** Persist the current state as one combined `{ messages, resume? }` record. */
+  private writeState(): void {
+    const messages = [...this.lastMessages]
+    // Nothing to persist (no transcript, no resume pointer): remove the key
+    // rather than writing an empty `{ messages: [] }`, so a cleared
+    // conversation does not leave a stale record behind.
+    if (messages.length === 0 && !this.lastResume) {
+      const generation = this.generation
+      this.runOperation(() => {
+        if (generation !== this.generation) {
+          return
+        }
+        return this.adapter.removeItem(this.id)
+      })
+      return
+    }
+    const generation = this.generation
+    const state: ChatPersistedState = {
+      messages,
+      ...(this.lastResume ? { resume: this.lastResume } : {}),
+    }
+    this.runOperation(() => {
+      if (generation !== this.generation) {
+        return
+      }
+      return this.adapter.setItem(this.id, state)
+    })
+  }
+
   // ---------------------------------------------------------------------------
   // Storage orchestration
   // ---------------------------------------------------------------------------
 
   /**
-   * Synchronously read the persisted messages for constructor-time hydration.
-   * Returns the raw `getItem` result (which may be a promise for async stores).
+   * Synchronously read the persisted state for constructor-time hydration.
+   * Returns the normalized combined record, or a promise of it for async stores.
    */
   readInitial():
-    | Array<UIMessage>
-    | null
+    | ChatPersistedState
     | undefined
-    | Promise<Array<UIMessage> | null | undefined> {
+    | Promise<ChatPersistedState | undefined> {
     try {
-      return this.adapter.getItem(this.id)
+      const raw = this.adapter.getItem(this.id)
+      if (raw instanceof Promise) {
+        return raw.then(normalizePersistedState).catch(() => undefined)
+      }
+      const state = normalizePersistedState(raw)
+      if (state) {
+        this.lastMessages = state.messages
+        this.lastResume = state.resume ?? null
+      }
+      return state
     } catch {
       return undefined
     }
   }
 
   /**
-   * Apply messages from an async `getItem` once it resolves, unless the message
+   * Apply state from an async `getItem` once it resolves, unless the message
    * list has already changed since hydration began.
    */
   hydrateAsync(
-    persistedMessages:
-      | Array<UIMessage>
-      | null
+    persistedState:
+      | ChatPersistedState
       | undefined
-      | Promise<Array<UIMessage> | null | undefined>,
+      | Promise<ChatPersistedState | undefined>,
   ): void {
-    if (!(persistedMessages instanceof Promise)) {
+    if (!(persistedState instanceof Promise)) {
       return
     }
 
     const hydrationGeneration = this.messagesGeneration
-    persistedMessages
-      .then((messages) => {
-        if (
-          Array.isArray(messages) &&
-          this.messagesGeneration === hydrationGeneration
-        ) {
-          this.applyMessages(messages)
+    persistedState
+      .then((state) => {
+        if (!state || this.messagesGeneration !== hydrationGeneration) {
+          return
+        }
+        this.lastResume = state.resume ?? null
+        this.lastMessages = state.messages
+        this.applyMessages(state.messages)
+        if (state.resume && this.applyResume) {
+          this.applyResume(state.resume)
         }
       })
       .catch(() => {
@@ -116,28 +176,37 @@ export class ChatPersistor {
   }
 
   /**
-   * Record a message-list change and queue a `setItem` write for it. Skips a
+   * Record a message-list change and queue a combined write for it. Skips a
    * single write after {@link beginClear} so the clear's empty snapshot isn't
    * persisted between `clearMessages()` and {@link remove}.
    */
   notifyMessagesChanged(messages: Array<UIMessage>): void {
     this.messagesGeneration++
+    this.lastMessages = [...messages]
     if (this.skipNextPersist) {
       this.skipNextPersist = false
       return
     }
-    const generation = this.generation
-    const messagesSnapshot = [...messages]
-    this.runOperation(() => {
-      if (generation !== this.generation) {
-        return
-      }
-      return this.adapter.setItem(this.id, messagesSnapshot)
-    })
+    this.writeState()
+  }
+
+  /**
+   * Record the current resume snapshot (which run to rejoin / which interrupts
+   * are pending) and persist it alongside the messages. Pass `null` to clear it
+   * once the run reaches a non-interrupt terminal.
+   */
+  persistResumeSnapshot(snapshot: ChatResumeSnapshot | null): void {
+    this.lastResume = snapshot
+    if (this.skipNextPersist) {
+      return
+    }
+    this.writeState()
   }
 
   /** Remove the persisted conversation. Invalidates any queued writes. */
   remove(): void {
+    this.lastMessages = []
+    this.lastResume = null
     const generation = ++this.generation
     this.runOperation(() => {
       if (generation !== this.generation) {