@tanstack/ai/src/activities/stream-generation-result.ts

Type-safe TypeScript AI SDK for streaming chat, tool calling, agents, structured outputs, and multimodal generation.

/**
 * Internal helper for wrapping one-shot generation results as StreamChunk
 * async iterables. NOT exported from the package — used only by activity
 * implementations to support `stream: true`.
 */

import { EventType } from '@ag-ui/core'
import { toRunErrorPayload } from './error-payload'
import type { StreamChunk } from '../types'

function createId(prefix: string): string {
  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
}

/**
 * Persisted artifact refs a middleware may have attached to the result. Read
 * defensively: the result shape is activity-specific and `artifacts` is only
 * present when generation persistence is wired with an artifact + blob store.
 */
function artifactsFromResult(result: unknown): Array<unknown> | undefined {
  if (typeof result !== 'object' || result === null) return undefined
  const artifacts = (result as { artifacts?: unknown }).artifacts
  return Array.isArray(artifacts) && artifacts.length > 0
    ? artifacts
    : undefined
}

/**
 * Wrap a one-shot generation result as a StreamChunk async iterable.
 *
 * This allows non-streaming activities (image, speech, transcription, summarize)
 * to be sent over the same streaming transport as chat.
 *
 * @param generator - An async function that performs the generation and returns the result
 * @param options - Optional configuration (runId, threadId)
 * @returns An AsyncIterable of StreamChunks with RUN_STARTED, CUSTOM(generation:result), and RUN_FINISHED events on success, or RUN_STARTED and RUN_ERROR on failure
 */
export async function* streamGenerationResult<TResult>(
  generator: (resolved: {
    runId: string
    threadId: string
  }) => Promise<TResult>,
  options?: { runId?: string; threadId?: string },
): AsyncIterable<StreamChunk> {
  const runId = options?.runId ?? createId('run')
  const threadId = options?.threadId ?? createId('thread')

  yield {
    type: EventType.RUN_STARTED,
    runId,
    threadId,
    timestamp: Date.now(),
  }

  try {
    const result = await generator({ runId, threadId })

    // Emit persisted artifact refs (if a middleware attached any) before the
    // result, so the client records them as the run streams.
    const artifacts = artifactsFromResult(result)
    if (artifacts) {
      yield {
        type: EventType.CUSTOM,
        name: 'generation:artifacts',
        value: artifacts,
        timestamp: Date.now(),
      }
    }

    yield {
      type: EventType.CUSTOM,
      name: 'generation:result',
      value: result as unknown,
      timestamp: Date.now(),
    }

    yield {
      type: EventType.RUN_FINISHED,
      runId,
      threadId,
      finishReason: 'stop',
      timestamp: Date.now(),
    }
  } catch (error: unknown) {
    const payload = toRunErrorPayload(error, 'Generation failed')
    // `code` is omitted entirely when undefined so the event matches the
    // AG-UI `code?: string` shape under `exactOptionalPropertyTypes`. The
    // deprecated nested `error` form preserves the same conditional
    // structure for backward compatibility.
    const codeFields =
      payload.code !== undefined ? { code: payload.code } : undefined
    yield {
      type: EventType.RUN_ERROR,
      message: payload.message,
      ...codeFields,
      // Deprecated nested form for backward compatibility
      error: {
        message: payload.message,
        ...codeFields,
      },
      timestamp: Date.now(),
    }
  }
}