@tanstack/ai 0.6.3 → 0.8.1

package/dist/esm/activities/chat/middleware/types.d.ts

@@ -0,0 +1,307 @@
+import { ModelMessage, StreamChunk, Tool, ToolCall } from '../../../types.js';
+/**
+ * Phase of the chat middleware lifecycle.
+ * - 'init': Initial config transform before the chat engine starts
+ * - 'beforeModel': Before each adapter chatStream call (per agent iteration)
+ * - 'modelStream': During model streaming
+ * - 'beforeTools': Before tool execution phase
+ * - 'afterTools': After tool execution phase
+ */
+export type ChatMiddlewarePhase = 'init' | 'beforeModel' | 'modelStream' | 'beforeTools' | 'afterTools';
+/**
+ * Stable context object passed to all middleware hooks.
+ * Created once per chat() invocation and shared across all hooks.
+ */
+export interface ChatMiddlewareContext {
+    /** Unique identifier for this chat request */
+    requestId: string;
+    /** Unique identifier for this stream */
+    streamId: string;
+    /** Conversation identifier, if provided by the caller */
+    conversationId?: string;
+    /** Current lifecycle phase */
+    phase: ChatMiddlewarePhase;
+    /** Current agent loop iteration (0-indexed) */
+    iteration: number;
+    /** Running count of chunks yielded so far */
+    chunkIndex: number;
+    /** Abort signal from the chat request */
+    signal?: AbortSignal;
+    /** Abort the chat run with a reason */
+    abort: (reason?: string) => void;
+    /** Opaque user-provided value from chat() options */
+    context: unknown;
+    /**
+     * Defer a non-blocking side-effect promise.
+     * Deferred promises do not block streaming and are awaited
+     * after the terminal hook (onFinish/onAbort/onError).
+     */
+    defer: (promise: Promise<unknown>) => void;
+    /** Provider name (e.g., 'openai', 'anthropic') */
+    provider: string;
+    /** Model identifier (e.g., 'gpt-4o') */
+    model: string;
+    /** Source of the chat invocation — always 'server' for server-side chat */
+    source: 'client' | 'server';
+    /** Whether the chat is streaming */
+    streaming: boolean;
+    /** System prompts configured for this chat */
+    systemPrompts: Array<string>;
+    /** Names of configured tools, if any */
+    toolNames?: Array<string>;
+    /** Flattened generation options (temperature, topP, maxTokens, metadata) */
+    options?: Record<string, unknown>;
+    /** Provider-specific model options */
+    modelOptions?: Record<string, unknown>;
+    /** Number of messages at the start of the request */
+    messageCount: number;
+    /** Whether tools are configured */
+    hasTools: boolean;
+    /** Current assistant message ID (changes per iteration) */
+    currentMessageId: string | null;
+    /** Accumulated text content for the current iteration */
+    accumulatedContent: string;
+    /** Current messages array (read-only view) */
+    messages: ReadonlyArray<ModelMessage>;
+    /** Generate a unique ID with the given prefix */
+    createId: (prefix: string) => string;
+}
+/**
+ * Chat configuration that middleware can observe or transform.
+ * This is a subset of the chat engine's effective configuration
+ * that middleware is allowed to modify.
+ */
+export interface ChatMiddlewareConfig {
+    messages: Array<ModelMessage>;
+    systemPrompts: Array<string>;
+    tools: Array<Tool>;
+    temperature?: number;
+    topP?: number;
+    maxTokens?: number;
+    metadata?: Record<string, unknown>;
+    modelOptions?: Record<string, unknown>;
+}
+/**
+ * Context provided to tool call hooks (onBeforeToolCall / onAfterToolCall).
+ */
+export interface ToolCallHookContext {
+    /** The tool call being executed */
+    toolCall: ToolCall;
+    /** The resolved tool definition, if found */
+    tool: Tool | undefined;
+    /** Parsed arguments for the tool call */
+    args: unknown;
+    /** Name of the tool */
+    toolName: string;
+    /** ID of the tool call */
+    toolCallId: string;
+}
+/**
+ * Decision returned from onBeforeToolCall.
+ * - undefined/void: continue with normal execution
+ * - { type: 'transformArgs', args }: replace args used for execution
+ * - { type: 'skip', result }: skip execution, use provided result
+ * - { type: 'abort', reason }: abort the entire chat run
+ */
+export type BeforeToolCallDecision = void | undefined | null | {
+    type: 'transformArgs';
+    args: unknown;
+} | {
+    type: 'skip';
+    result: unknown;
+} | {
+    type: 'abort';
+    reason?: string;
+};
+/**
+ * Outcome information provided to onAfterToolCall.
+ */
+export interface AfterToolCallInfo {
+    /** The tool call that was executed */
+    toolCall: ToolCall;
+    /** The resolved tool definition */
+    tool: Tool | undefined;
+    /** Name of the tool */
+    toolName: string;
+    /** ID of the tool call */
+    toolCallId: string;
+    /** Whether the execution succeeded */
+    ok: boolean;
+    /** Duration of tool execution in milliseconds */
+    duration: number;
+    /** The result (if ok) or error (if not ok) */
+    result?: unknown;
+    error?: unknown;
+}
+/**
+ * Information passed to onIteration at the start of each agent loop iteration.
+ */
+export interface IterationInfo {
+    /** 0-based iteration index */
+    iteration: number;
+    /** The assistant message ID created for this iteration */
+    messageId: string;
+}
+/**
+ * Aggregate information passed to onToolPhaseComplete after all tool calls
+ * in an iteration have been processed.
+ */
+export interface ToolPhaseCompleteInfo {
+    /** Tool calls that were assigned to the assistant message */
+    toolCalls: Array<ToolCall>;
+    /** Completed tool results */
+    results: Array<{
+        toolCallId: string;
+        toolName: string;
+        result: unknown;
+        duration?: number;
+    }>;
+    /** Tools that need user approval */
+    needsApproval: Array<{
+        toolCallId: string;
+        toolName: string;
+        input: unknown;
+        approvalId: string;
+    }>;
+    /** Tools that need client-side execution */
+    needsClientExecution: Array<{
+        toolCallId: string;
+        toolName: string;
+        input: unknown;
+    }>;
+}
+/**
+ * Token usage statistics passed to the onUsage hook.
+ * Extracted from the RUN_FINISHED chunk when usage data is present.
+ */
+export interface UsageInfo {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+}
+/**
+ * Information passed to onFinish.
+ */
+export interface FinishInfo {
+    /** The finish reason from the last model response */
+    finishReason: string | null;
+    /** Total duration of the chat run in milliseconds */
+    duration: number;
+    /** Final accumulated text content */
+    content: string;
+    /** Final usage totals, if available */
+    usage?: {
+        promptTokens: number;
+        completionTokens: number;
+        totalTokens: number;
+    };
+}
+/**
+ * Information passed to onAbort.
+ */
+export interface AbortInfo {
+    /** The reason for the abort, if provided */
+    reason?: string;
+    /** Duration until abort in milliseconds */
+    duration: number;
+}
+/**
+ * Information passed to onError.
+ */
+export interface ErrorInfo {
+    /** The error that caused the failure */
+    error: unknown;
+    /** Duration until error in milliseconds */
+    duration: number;
+}
+/**
+ * Chat middleware interface.
+ *
+ * All hooks are optional. Middleware is composed in array order:
+ * - `onConfig`: config piped through middlewares in order (first transform influences later)
+ * - `onChunk`: each output chunk is fed into the next middleware in order
+ *
+ * @example Logging middleware
+ * ```ts
+ * const loggingMiddleware: ChatMiddleware = {
+ *   name: 'logging',
+ *   onStart(ctx) { console.log('Chat started', ctx.requestId) },
+ *   onChunk(ctx, chunk) { console.log('Chunk:', chunk.type) },
+ *   onFinish(ctx, info) { console.log('Done:', info.duration, 'ms') },
+ * }
+ * ```
+ *
+ * @example Redaction middleware
+ * ```ts
+ * const redactionMiddleware: ChatMiddleware = {
+ *   name: 'redaction',
+ *   onChunk(ctx, chunk) {
+ *     if (chunk.type === 'TEXT_MESSAGE_CONTENT') {
+ *       return { ...chunk, delta: redact(chunk.delta) }
+ *     }
+ *   },
+ * }
+ * ```
+ */
+export interface ChatMiddleware {
+    /** Optional name for debugging and identification */
+    name?: string;
+    /**
+     * Called to observe or transform the chat configuration.
+     * Called at init and at the beginning of each agent iteration.
+     *
+     * Return a partial config to merge with the current config, or void to pass through.
+     * Only the fields you return are overwritten — everything else is preserved.
+     */
+    onConfig?: (ctx: ChatMiddlewareContext, config: ChatMiddlewareConfig) => void | null | Partial<ChatMiddlewareConfig> | Promise<void | Partial<ChatMiddlewareConfig>>;
+    /**
+     * Called when the chat run starts (after initial onConfig).
+     */
+    onStart?: (ctx: ChatMiddlewareContext) => void | Promise<void>;
+    /**
+     * Called at the start of each agent loop iteration, after a new assistant message ID
+     * is created. Use this to observe iteration boundaries.
+     */
+    onIteration?: (ctx: ChatMiddlewareContext, info: IterationInfo) => void | Promise<void>;
+    /**
+     * Called for every chunk yielded by chat().
+     * Can observe, transform, expand, or drop chunks.
+     *
+     * @returns void (pass through), chunk (replace), chunk[] (expand), null (drop)
+     */
+    onChunk?: (ctx: ChatMiddlewareContext, chunk: StreamChunk) => void | StreamChunk | Array<StreamChunk> | null | Promise<void | StreamChunk | Array<StreamChunk> | null>;
+    /**
+     * Called before a tool is executed.
+     * Can observe, transform args, skip execution, or abort the run.
+     */
+    onBeforeToolCall?: (ctx: ChatMiddlewareContext, hookCtx: ToolCallHookContext) => BeforeToolCallDecision | Promise<BeforeToolCallDecision>;
+    /**
+     * Called after a tool execution completes (success or failure).
+     */
+    onAfterToolCall?: (ctx: ChatMiddlewareContext, info: AfterToolCallInfo) => void | Promise<void>;
+    /**
+     * Called after all tool calls in an iteration have been processed.
+     * Provides aggregate data about tool execution results, approvals, and client tools.
+     */
+    onToolPhaseComplete?: (ctx: ChatMiddlewareContext, info: ToolPhaseCompleteInfo) => void | Promise<void>;
+    /**
+     * Called when usage data is available from a RUN_FINISHED chunk.
+     * Called once per model iteration that reports usage.
+     */
+    onUsage?: (ctx: ChatMiddlewareContext, usage: UsageInfo) => void | Promise<void>;
+    /**
+     * Called when the chat run completes normally.
+     * Exactly one of onFinish/onAbort/onError will be called per run.
+     */
+    onFinish?: (ctx: ChatMiddlewareContext, info: FinishInfo) => void | Promise<void>;
+    /**
+     * Called when the chat run is aborted.
+     * Exactly one of onFinish/onAbort/onError will be called per run.
+     */
+    onAbort?: (ctx: ChatMiddlewareContext, info: AbortInfo) => void | Promise<void>;
+    /**
+     * Called when the chat run encounters an unhandled error.
+     * Exactly one of onFinish/onAbort/onError will be called per run.
+     */
+    onError?: (ctx: ChatMiddlewareContext, info: ErrorInfo) => void | Promise<void>;
+}

package/dist/esm/activities/chat/stream/processor.d.ts

@@ -50,9 +50,8 @@ export interface StreamProcessorOptions {
  *
  * State tracking:
  * - Full message array
- * - Current assistant message being streamed
- * - Text content accumulation (reset on TEXT_MESSAGE_START)
- * - Multiple parallel tool calls
+ * - Per-message stream state (text, tool calls, thinking)
+ * - Multiple concurrent message streams
  * - Tool call completion via TOOL_CALL_END events
  *
  * @see docs/chat-architecture.md#streamprocessor-internal-state — State field reference
@@ -64,13 +63,11 @@ export declare class StreamProcessor {
     private jsonParser;
     private recordingEnabled;
     private messages;
-    private currentAssistantMessageId;
-    private totalTextContent;
-    private currentSegmentText;
-    private lastEmittedText;
-    private thinkingContent;
-    private toolCalls;
-    private toolCallOrder;
+    private messageStates;
+    private activeMessageIds;
+    private toolCallToMessage;
+    private pendingManualMessageId;
+    private activeRuns;
     private finishReason;
     private hasError;
     private isDone;
@@ -117,24 +114,13 @@ export declare class StreamProcessor {
      * @deprecated Use prepareAssistantMessage() instead. This eagerly creates
      * an assistant message which can cause empty message flicker.
      */
-    startAssistantMessage(): string;
+    startAssistantMessage(messageId?: string): string;
     /**
      * Get the current assistant message ID (if one has been created).
      * Returns null if prepareAssistantMessage() was called but no content
      * has arrived yet.
      */
     getCurrentAssistantMessageId(): string | null;
-    /**
-     * Lazily create the assistant message if it hasn't been created yet.
-     * Called by content handlers on the first content-bearing chunk.
-     * Returns the message ID.
-     *
-     * Content-bearing chunks that trigger this:
-     * TEXT_MESSAGE_CONTENT, TOOL_CALL_START, STEP_FINISHED, RUN_ERROR.
-     *
-     * @see docs/chat-architecture.md#streamprocessor-internal-state — Lazy creation pattern
-     */
-    private ensureAssistantMessage;
     /**
      * Add a tool result (called by client after handling onToolCall)
      */
@@ -173,23 +159,46 @@ export declare class StreamProcessor {
      *
      * Central dispatch for all AG-UI events. Each event type maps to a specific
      * handler. Events not listed in the switch are intentionally ignored
-     * (RUN_STARTED, TEXT_MESSAGE_END, STEP_STARTED, STATE_SNAPSHOT, STATE_DELTA).
+     * (RUN_STARTED, STEP_STARTED, STATE_DELTA).
      *
      * @see docs/chat-architecture.md#adapter-contract — Expected event types and ordering
      */
     processChunk(chunk: StreamChunk): void;
     /**
-     * Handle TEXT_MESSAGE_START event — marks the beginning of a new text segment.
-     * Resets segment accumulation so text after tool calls starts fresh.
-     *
-     * This is the key mechanism for multi-segment text (text before and after tool
-     * calls becoming separate TextParts). Without this reset, all text would merge
-     * into a single TextPart and tool-call interleaving would be lost.
+     * Create a new MessageStreamState for a message
+     */
+    private createMessageState;
+    /**
+     * Get the MessageStreamState for a message
+     */
+    private getMessageState;
+    /**
+     * Get the most recent active assistant message ID.
+     * Used as fallback for events that don't include a messageId.
+     */
+    private getActiveAssistantMessageId;
+    /**
+     * Ensure an active assistant message exists, creating one if needed.
+     * Used for backward compat when events arrive without prior TEXT_MESSAGE_START.
      *
-     * @see docs/chat-architecture.md#single-shot-text-response — Step-by-step text processing
-     * @see docs/chat-architecture.md#text-then-tool-interleaving-single-shot — Multi-segment text
+     * On reconnect/resume, a TEXT_MESSAGE_CONTENT may arrive for a message that
+     * already exists in this.messages (e.g. from initialMessages or a prior
+     * MESSAGES_SNAPSHOT) but whose transient state was cleared. In that case we
+     * hydrate state from the existing message rather than creating a duplicate.
+     */
+    private ensureAssistantMessage;
+    /**
+     * Handle TEXT_MESSAGE_START event
      */
     private handleTextMessageStartEvent;
+    /**
+     * Handle TEXT_MESSAGE_END event
+     */
+    private handleTextMessageEndEvent;
+    /**
+     * Handle MESSAGES_SNAPSHOT event
+     */
+    private handleMessagesSnapshotEvent;
     /**
      * Handle TEXT_MESSAGE_CONTENT event.
      *
@@ -244,12 +253,19 @@ export declare class StreamProcessor {
      * @see docs/chat-architecture.md#single-shot-tool-call-response — End-to-end flow
      */
     private handleToolCallEndEvent;
+    /**
+     * Handle RUN_STARTED event.
+     *
+     * Registers the run so that RUN_FINISHED can determine whether other
+     * runs are still active before finalizing.
+     */
+    private handleRunStartedEvent;
     /**
      * Handle RUN_FINISHED event.
      *
-     * Records the finishReason and calls completeAllToolCalls() as a safety net
-     * to force-complete any tool calls that didn't receive an explicit TOOL_CALL_END.
-     * This handles cases like aborted streams or adapter bugs.
+     * Records the finishReason and removes the run from activeRuns.
+     * Only finalizes when no more runs are active, so that concurrent
+     * runs don't interfere with each other.
      *
      * @see docs/chat-architecture.md#single-shot-tool-call-response — finishReason semantics
      * @see docs/chat-architecture.md#adapter-contract — Why RUN_FINISHED is mandatory
@@ -280,7 +296,11 @@ export declare class StreamProcessor {
      */
     private handleCustomEvent;
     /**
-     * Complete all tool calls — safety net for stream termination.
+     * Detect if an incoming content chunk represents a NEW text segment
+     */
+    private isNewTextSegment;
+    /**
+     * Complete all tool calls across all active messages — safety net for stream termination.
      *
      * Called by RUN_FINISHED and finalizeStream(). Force-transitions any tool call
      * not yet in input-complete state. Handles cases where TOOL_CALL_END was
@@ -289,12 +309,16 @@ export declare class StreamProcessor {
      * @see docs/chat-architecture.md#single-shot-tool-call-response — Safety net behavior
      */
     private completeAllToolCalls;
+    /**
+     * Complete all tool calls for a specific message
+     */
+    private completeAllToolCallsForMessage;
     /**
      * Mark a tool call as complete and emit event
      */
     private completeToolCall;
     /**
-     * Emit pending text update.
+     * Emit pending text update for a specific message.
      *
      * Calls updateTextPart() which has critical append-vs-replace logic:
      * - If last UIMessage part is TextPart → replaces its content (same segment).
@@ -302,7 +326,7 @@ export declare class StreamProcessor {
      *
      * @see docs/chat-architecture.md#uimessage-part-ordering-invariants — Replace vs. push logic
      */
-    private emitTextUpdate;
+    private emitTextUpdateForMessage;
     /**
      * Emit messages change event
      */
@@ -318,15 +342,15 @@ export declare class StreamProcessor {
      */
     finalizeStream(): void;
     /**
-     * Get completed tool calls in API format
+     * Get completed tool calls in API format (aggregated across all messages)
      */
     private getCompletedToolCalls;
     /**
-     * Get current result
+     * Get current result (aggregated across all messages)
      */
     private getResult;
     /**
-     * Get current processor state
+     * Get current processor state (aggregated across all messages)
      */
     getState(): ProcessorState;
     /**