assistant-cloud/src/runTelemetry.ts

Cloud integration for assistant-ui

import type { SamplingCallData } from "./instrumentMcpSampling";
import type { AssistantCloudRunReport } from "./AssistantCloudRuns";

const MAX_TELEMETRY_TEXT_LENGTH = 50_000;

const BASE64_PATTERN = /^[A-Za-z0-9+/]+={0,2}$/;

export type AssistantCloudRunReportToolCall = {
  tool_name: string;
  tool_call_id: string;
  tool_args?: string;
  tool_result?: string;
  tool_source?: "mcp" | "frontend" | "backend";
  start_ms?: number;
  end_ms?: number;
  sampling_calls?: SamplingCallData[];
};

export type RunReportOutcome =
  | "aborted"
  | "disconnected"
  | "length"
  | "content_filter"
  | "rate_limited"
  | "validation_failed"
  | "server_error"
  | "persistence_error";

/**
 * Maps a finish event to the report status and outcome. `fallbackStatus`
 * applies when the event carries neither a finish reason nor a failure flag.
 */
export function deriveRunOutcome(
  input: {
    finishReason?: string | undefined;
    isAbort?: boolean | undefined;
    isDisconnect?: boolean | undefined;
    isError?: boolean | undefined;
  },
  fallbackStatus: "completed" | "incomplete" = "completed",
): {
  status: "completed" | "incomplete" | "error";
  outcome?: RunReportOutcome;
} {
  if (input.isError) return { status: "error" };
  if (input.isAbort) return { status: "incomplete", outcome: "aborted" };
  if (input.isDisconnect) {
    return { status: "incomplete", outcome: "disconnected" };
  }
  switch (input.finishReason) {
    case "length":
      return { status: "incomplete", outcome: "length" };
    case "content-filter":
    case "content_filter":
      return { status: "incomplete", outcome: "content_filter" };
    case "cancelled":
      return { status: "incomplete", outcome: "aborted" };
    case "error":
      return { status: "error" };
    default:
      return {
        status: input.finishReason === undefined ? fallbackStatus : "completed",
      };
  }
}

const MAX_RUN_ERROR_CODE_LENGTH = 64;
const MAX_RUN_ERROR_LENGTH = 2048;

/**
 * Reads the message, code and outcome the runs endpoint stores for a failed
 * run. The code is the error's HTTP `statusCode` when it carries one, else its
 * `code`, else its class name. The outcome comes from the HTTP status.
 */
export function describeRunError(error: unknown): {
  error?: string;
  errorCode?: string;
  outcome?: RunReportOutcome;
} {
  if (error == null) return {};
  const record =
    typeof error === "object" ? (error as Record<string, unknown>) : undefined;
  const message =
    typeof error === "string"
      ? error
      : typeof record?.message === "string"
        ? record.message
        : undefined;
  const statusCode =
    typeof record?.statusCode === "number" &&
    Number.isInteger(record.statusCode)
      ? record.statusCode
      : undefined;
  const code =
    statusCode !== undefined
      ? String(statusCode)
      : typeof record?.code === "string"
        ? record.code
        : typeof record?.name === "string" && record.name !== "Error"
          ? record.name
          : undefined;
  const outcome =
    statusCode !== undefined ? outcomeForStatusCode(statusCode) : undefined;
  return {
    ...(message
      ? { error: message.slice(0, MAX_RUN_ERROR_LENGTH) }
      : undefined),
    ...(code
      ? { errorCode: code.slice(0, MAX_RUN_ERROR_CODE_LENGTH) }
      : undefined),
    ...(outcome ? { outcome } : undefined),
  };
}

const outcomeForStatusCode = (
  statusCode: number,
): RunReportOutcome | undefined => {
  if (statusCode === 429) return "rate_limited";
  if (statusCode === 400 || statusCode === 422) return "validation_failed";
  if (statusCode >= 500 && statusCode < 600) return "server_error";
  return undefined;
};

/**
 * Clamps a string to the size the runs endpoint accepts for a single span
 * field.
 */
export function truncateRunTelemetryText(value: string): string {
  if (value.length <= MAX_TELEMETRY_TEXT_LENGTH) return value;
  return value.slice(0, MAX_TELEMETRY_TEXT_LENGTH);
}

function safeStringify(value: unknown): string | undefined {
  if (value == null) return undefined;
  try {
    return truncateRunTelemetryText(JSON.stringify(value));
  } catch {
    return undefined;
  }
}

const base64SizeKB = (value: string) =>
  ((value.length * 3) / 4 / 1024).toFixed(1);

// Both call sites read a field the MCP content grammar already defines as
// base64, so the test only has to recognise the encoding's own shape: its
// alphabet, and a length that is a multiple of four. A size floor would leave
// a short payload — `aGk=` is a whole image in this repo's own fixtures —
// serialized raw.
const isInlineBase64 = (value: unknown): value is string => {
  if (typeof value !== "string") return false;
  const head = value.slice(0, 200);
  return head.length > 0 && head.length % 4 === 0 && BASE64_PATTERN.test(head);
};

function summarizeMcpContentBlock(item: unknown): unknown {
  if (!item || typeof item !== "object") return item;
  const block = item as {
    type?: unknown;
    data?: unknown;
    resource?: unknown;
  };
  if (
    (block.type === "image" || block.type === "audio") &&
    isInlineBase64(block.data)
  ) {
    return {
      ...block,
      data: `[${String(block.type)}: ${base64SizeKB(block.data)}KB]`,
    };
  }
  // An EmbeddedResource is the third inline carrier: a binary resource arrives
  // as base64 under `resource.blob` rather than as a top-level `data` field.
  if (
    block.type === "resource" &&
    block.resource &&
    typeof block.resource === "object"
  ) {
    const resource = block.resource as { blob?: unknown };
    if (isInlineBase64(resource.blob)) {
      return {
        ...block,
        resource: {
          ...resource,
          blob: `[resource: ${base64SizeKB(resource.blob)}KB]`,
        },
      };
    }
  }
  return item;
}

function summarizeMcpResult(value: unknown): string | undefined {
  if (value == null) return undefined;
  try {
    const parsed = typeof value === "string" ? JSON.parse(value) : value;
    if (Array.isArray(parsed)) {
      return truncateRunTelemetryText(
        JSON.stringify(parsed.map(summarizeMcpContentBlock)),
      );
    }
    // `callTool` resolves to a CallToolResult, so the blocks carrying base64
    // arrive under `content` rather than as the result itself.
    if (typeof parsed === "object") {
      const content = (parsed as { content?: unknown }).content;
      if (Array.isArray(content)) {
        return truncateRunTelemetryText(
          JSON.stringify({
            ...(parsed as Record<string, unknown>),
            content: content.map(summarizeMcpContentBlock),
          }),
        );
      }
    }
  } catch {
    // not a shape with summarizable content, fall through
  }
  return safeStringify(value);
}

export type RunTelemetryToolCallInit = {
  toolName: string;
  toolCallId: string;
  args?: unknown;
  /**
   * Pre-serialized arguments, used in place of serializing `args`. Values over
   * the span size are clamped before they are included in the report.
   */
  argsText?: string | undefined;
  result?: unknown;
  toolSource?: "mcp" | "frontend" | "backend" | undefined;
};

/**
 * Serializes one tool call into the shape the runs endpoint accepts. An `mcp`
 * source has its result summarized, because MCP content blocks carry inline
 * base64 image and audio payloads that would otherwise dominate the report.
 */
export function createRunTelemetryToolCall(
  init: RunTelemetryToolCallInit,
): AssistantCloudRunReportToolCall {
  const { toolName, toolCallId, args, argsText, result, toolSource } = init;
  const call: AssistantCloudRunReportToolCall = {
    tool_name: toolName,
    tool_call_id: toolCallId,
  };
  const toolArgs =
    argsText != null ? truncateRunTelemetryText(argsText) : safeStringify(args);
  if (toolArgs !== undefined) call.tool_args = toolArgs;
  const toolResult =
    toolSource === "mcp" ? summarizeMcpResult(result) : safeStringify(result);
  if (toolResult !== undefined) call.tool_result = toolResult;
  if (toolSource) call.tool_source = toolSource;
  return call;
}

/**
 * Resolves the model ID a run reports, in the order an app can supply it: an
 * explicit `modelId`, the `custom` bag, then the per-step `response.modelId`
 * that a `messageMetadata` callback copies off the AI SDK's finish-step part.
 * The AI SDK puts no model ID on a UI message part, so message metadata is the
 * only channel one arrives on.
 */
export function extractRunTelemetryModelId(
  metadata: Record<string, unknown> | undefined,
): string | undefined {
  if (!metadata) return undefined;
  if (typeof metadata.modelId === "string") return metadata.modelId;
  const custom = metadata.custom as Record<string, unknown> | undefined;
  if (typeof custom?.modelId === "string") return custom.modelId;

  const steps: unknown = metadata.steps;
  if (!Array.isArray(steps)) return undefined;
  for (const step of steps as unknown[]) {
    if (!step || typeof step !== "object") continue;
    const response = (step as Record<string, unknown>).response;
    if (!response || typeof response !== "object") continue;
    const modelId = (response as Record<string, unknown>).modelId;
    if (typeof modelId === "string") return modelId;
  }
  return undefined;
}

export type RunTelemetryUsage = {
  inputTokens?: number | undefined;
  outputTokens?: number | undefined;
  reasoningTokens?: number | undefined;
  cachedInputTokens?: number | undefined;
};

export type RunTelemetryUsageInit = RunTelemetryUsage & {
  promptTokens?: number | undefined;
  completionTokens?: number | undefined;
  inputTokenDetails?: { cacheReadTokens?: number };
  outputTokenDetails?: { reasoningTokens?: number };
};

export type RunReportStepInit = {
  usage?: RunTelemetryUsageInit | undefined;
  toolCalls?: AssistantCloudRunReportToolCall[] | undefined;
  startMs?: number | undefined;
  endMs?: number | undefined;
  finishReason?: string | undefined;
  input?: string | undefined;
};

/**
 * The run report fields read from the messages of one run, in whichever
 * format they were stored: the status the messages imply, the tool calls, the
 * steps, the text, the usage and the model.
 */
export type RunMessageTelemetry = {
  assistantMessageId?: string;
  status: "completed" | "incomplete";
  toolCalls?: AssistantCloudRunReportToolCall[];
  steps?: RunReportStepInit[];
  totalSteps?: number;
  outputText?: string;
  usage?: RunTelemetryUsage;
  modelId?: string;
  metadata?: Record<string, unknown>;
};

export type RunReportInit = {
  threadId: string;
  status: AssistantCloudRunReport["status"];
  outcome?: AssistantCloudRunReport["outcome_type"] | undefined;
  errorCode?: string | undefined;
  error?: string | undefined;
  messageId?: string | undefined;
  userMessageId?: string | undefined;
  traceId?: string | undefined;
  rootSpanId?: string | undefined;
  modelId?: string | undefined;
  provider?: string | undefined;
  usage?: RunTelemetryUsageInit | undefined;
  steps?: RunReportStepInit[] | undefined;
  totalSteps?: number | undefined;
  toolCalls?: AssistantCloudRunReportToolCall[] | undefined;
  durationMs?: number | undefined;
  firstTokenMs?: number | undefined;
  costUsd?: number | undefined;
  costDetails?:
    | {
        input?: number | undefined;
        inputCachedTokens?: number | undefined;
        output?: number | undefined;
        total?: number | undefined;
      }
    | undefined;
  outputText?: string | undefined;
  attributes?: Record<string, unknown> | undefined;
  metadata?: Record<string, unknown> | undefined;
  telemetry?: {
    environment?: string | undefined;
    release?: string | undefined;
    tags?: readonly string[] | undefined;
  };
};

function assignUsage(
  report: AssistantCloudRunReport,
  usage: RunTelemetryUsageInit | undefined,
): void {
  if (!usage) return;
  const normalized = normalizeRunTelemetryUsage(usage);
  if (!normalized) return;
  if (normalized.inputTokens !== undefined) {
    report.input_tokens = normalized.inputTokens;
  }
  if (normalized.outputTokens !== undefined) {
    report.output_tokens = normalized.outputTokens;
  }
  if (normalized.reasoningTokens !== undefined) {
    report.reasoning_tokens = normalized.reasoningTokens;
  }
  if (normalized.cachedInputTokens !== undefined) {
    report.cached_input_tokens = normalized.cachedInputTokens;
  }
}

function createRunReportStep(
  init: RunReportStepInit,
): NonNullable<AssistantCloudRunReport["steps"]>[number] {
  const step: NonNullable<AssistantCloudRunReport["steps"]>[number] = {};
  const usage = init.usage ? normalizeRunTelemetryUsage(init.usage) : undefined;
  if (usage?.inputTokens !== undefined) step.input_tokens = usage.inputTokens;
  if (usage?.outputTokens !== undefined)
    step.output_tokens = usage.outputTokens;
  if (usage?.reasoningTokens !== undefined) {
    step.reasoning_tokens = usage.reasoningTokens;
  }
  if (usage?.cachedInputTokens !== undefined) {
    step.cached_input_tokens = usage.cachedInputTokens;
  }
  if (init.toolCalls !== undefined) step.tool_calls = init.toolCalls;
  if (init.startMs !== undefined) step.start_ms = init.startMs;
  if (init.endMs !== undefined) step.end_ms = init.endMs;
  if (init.finishReason !== undefined) {
    step.finish_reason = init.finishReason.slice(0, 32);
  }
  if (init.input !== undefined) {
    step.input = truncateRunTelemetryText(init.input);
  }
  return step;
}

function normalizeRunReportTags(
  tags: readonly string[] | undefined,
): string[] | undefined {
  if (!tags) return undefined;
  const result: string[] = [];
  const seen = new Set<string>();
  for (const tag of tags) {
    const normalized = tag.trim().slice(0, 64).trim();
    if (!normalized || seen.has(normalized)) continue;
    seen.add(normalized);
    result.push(normalized);
    if (result.length === 20) break;
  }
  return result.length > 0 ? result : undefined;
}

function normalizeRunReportMilliseconds(
  value: number | undefined,
): number | undefined {
  if (value === undefined || !Number.isFinite(value)) return undefined;
  return Math.max(0, Math.round(value));
}

function normalizeRunReportCost(value: number | undefined): number | undefined {
  if (value === undefined || !Number.isFinite(value) || value < 0) {
    return undefined;
  }
  return value;
}

function assignCostDetails(
  report: AssistantCloudRunReport,
  costDetails: RunReportInit["costDetails"],
): void {
  if (!costDetails) return;
  const normalized: NonNullable<AssistantCloudRunReport["cost_details"]> = {};
  const input = normalizeRunReportCost(costDetails.input);
  if (input !== undefined) normalized.input = input;
  const inputCachedTokens = normalizeRunReportCost(
    costDetails.inputCachedTokens,
  );
  if (inputCachedTokens !== undefined) {
    normalized.input_cached_tokens = inputCachedTokens;
  }
  const output = normalizeRunReportCost(costDetails.output);
  if (output !== undefined) normalized.output = output;
  const total = normalizeRunReportCost(costDetails.total);
  if (total !== undefined) normalized.total = total;
  if (Object.keys(normalized).length > 0) report.cost_details = normalized;
}

export function createRunReport(init: RunReportInit): AssistantCloudRunReport {
  const report: AssistantCloudRunReport = {
    thread_id: init.threadId,
    status: init.status,
  };
  const traceId = init.traceId?.toLowerCase();
  if (traceId && /^[0-9a-f]{32}$/.test(traceId)) report.trace_id = traceId;
  const rootSpanId = init.rootSpanId?.toLowerCase();
  if (rootSpanId && /^[0-9a-f]{16}$/.test(rootSpanId)) {
    report.root_span_id = rootSpanId;
  }
  if (init.outcome !== undefined) report.outcome_type = init.outcome;
  if (init.errorCode !== undefined) report.error_code = init.errorCode;
  if (init.error !== undefined) report.error = init.error;
  if (init.messageId !== undefined) report.message_id = init.messageId;
  if (init.userMessageId !== undefined)
    report.user_message_id = init.userMessageId;
  if (init.modelId !== undefined) report.model_id = init.modelId;
  if (init.provider !== undefined) {
    report.provider = init.provider;
    report.provider_type = init.provider;
  }
  assignUsage(report, init.usage);
  if (init.steps !== undefined) {
    report.steps = init.steps.map(createRunReportStep);
    report.total_steps = init.steps.length;
  } else if (init.totalSteps !== undefined) {
    report.total_steps = init.totalSteps;
  }
  if (init.toolCalls !== undefined) report.tool_calls = init.toolCalls;
  const durationMs = normalizeRunReportMilliseconds(init.durationMs);
  if (durationMs !== undefined) report.duration_ms = durationMs;
  const firstTokenMs = normalizeRunReportMilliseconds(init.firstTokenMs);
  if (firstTokenMs !== undefined) report.first_token_ms = firstTokenMs;
  const costUsd = normalizeRunReportCost(init.costUsd);
  if (costUsd !== undefined) report.cost_usd = costUsd;
  assignCostDetails(report, init.costDetails);
  if (init.outputText !== undefined) {
    report.output_text = truncateRunTelemetryText(init.outputText);
  }
  if (init.attributes !== undefined) report.attributes = init.attributes;
  if (init.metadata !== undefined) report.metadata = init.metadata;
  if (init.telemetry?.environment !== undefined) {
    report.environment = init.telemetry.environment;
  }
  if (init.telemetry?.release !== undefined) {
    report.release = init.telemetry.release;
  }
  const tags = normalizeRunReportTags(init.telemetry?.tags);
  if (tags !== undefined) report.tags = tags;
  return report;
}

/**
 * Resolves the token counts a provider reports under any of the names the AI
 * SDK has used: the current top-level ones, the legacy prompt/completion pair,
 * and the v7 token detail objects. Returns undefined when no count is present,
 * so callers can tell an empty usage object from a zeroed one.
 */
export function normalizeRunTelemetryUsage(
  usage: RunTelemetryUsageInit,
): RunTelemetryUsage | undefined {
  const inputTokens = usage.inputTokens ?? usage.promptTokens;
  const outputTokens = usage.outputTokens ?? usage.completionTokens;
  // AI SDK v7 moved these under token detail objects; v6 kept them top-level.
  const reasoningTokens =
    usage.reasoningTokens ?? usage.outputTokenDetails?.reasoningTokens;
  const cachedInputTokens =
    usage.cachedInputTokens ?? usage.inputTokenDetails?.cacheReadTokens;

  if (
    inputTokens == null &&
    outputTokens == null &&
    reasoningTokens == null &&
    cachedInputTokens == null
  ) {
    return undefined;
  }

  return {
    ...(inputTokens != null ? { inputTokens } : undefined),
    ...(outputTokens != null ? { outputTokens } : undefined),
    ...(reasoningTokens != null ? { reasoningTokens } : undefined),
    ...(cachedInputTokens != null ? { cachedInputTokens } : undefined),
  };
}