assistant-cloud 0.1.42 → 0.2.0

package/src/runTelemetry.ts

@@ -1,4 +1,5 @@
 import type { SamplingCallData } from "./instrumentMcpSampling";
+import type { AssistantCloudRunReport } from "./AssistantCloudRuns";
 
 const MAX_TELEMETRY_TEXT_LENGTH = 50_000;
 
@@ -15,6 +16,86 @@ export type AssistantCloudRunReportToolCall = {
   sampling_calls?: SamplingCallData[];
 };
 
+export type RunReportOutcome =
+  | "aborted"
+  | "disconnected"
+  | "length"
+  | "content_filter";
+
+/**
+ * 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 and code the runs endpoint stores for a failed run. The
+ * code is the error's `code` when it has one, else its class name.
+ */
+export function describeRunError(error: unknown): {
+  error?: string;
+  errorCode?: string;
+} {
+  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 code =
+    typeof record?.code === "string"
+      ? record.code
+      : typeof record?.name === "string" && record.name !== "Error"
+        ? record.name
+        : undefined;
+  return {
+    ...(message
+      ? { error: message.slice(0, MAX_RUN_ERROR_LENGTH) }
+      : undefined),
+    ...(code
+      ? { errorCode: code.slice(0, MAX_RUN_ERROR_CODE_LENGTH) }
+      : undefined),
+  };
+}
+
 /**
  * Clamps a string to the size the runs endpoint accepts for a single span
  * field.
@@ -95,21 +176,194 @@ export function createRunTelemetryToolCall(
   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;
-  outputTokens?: number;
-  reasoningTokens?: number;
-  cachedInputTokens?: number;
+  inputTokens?: number | undefined;
+  outputTokens?: number | undefined;
+  reasoningTokens?: number | undefined;
+  cachedInputTokens?: number | undefined;
 };
 
 export type RunTelemetryUsageInit = RunTelemetryUsage & {
-  promptTokens?: number;
-  completionTokens?: number;
+  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;
+};
+
+export type RunReportInit = {
+  threadId: string;
+  status: AssistantCloudRunReport["status"];
+  outcome?: RunReportOutcome | undefined;
+  errorCode?: string | undefined;
+  error?: string | undefined;
+  messageId?: string | undefined;
+  traceId?: 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;
+  outputText?: string | 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);
+  }
+  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));
+}
+
+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;
+  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.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;
+  if (init.outputText !== undefined) {
+    report.output_text = truncateRunTelemetryText(init.outputText);
+  }
+  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 either the current or the
- * legacy prompt/completion names. Returns undefined when no count is present,
+ * 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(
@@ -117,12 +371,17 @@ export function normalizeRunTelemetryUsage(
 ): 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 &&
-    usage.reasoningTokens == null &&
-    usage.cachedInputTokens == null
+    reasoningTokens == null &&
+    cachedInputTokens == null
   ) {
     return undefined;
   }
@@ -130,11 +389,7 @@ export function normalizeRunTelemetryUsage(
   return {
     ...(inputTokens != null ? { inputTokens } : undefined),
     ...(outputTokens != null ? { outputTokens } : undefined),
-    ...(usage.reasoningTokens != null
-      ? { reasoningTokens: usage.reasoningTokens }
-      : undefined),
-    ...(usage.cachedInputTokens != null
-      ? { cachedInputTokens: usage.cachedInputTokens }
-      : undefined),
+    ...(reasoningTokens != null ? { reasoningTokens } : undefined),
+    ...(cachedInputTokens != null ? { cachedInputTokens } : undefined),
   };
 }

package/src/telemetry/index.test.ts

@@ -0,0 +1,252 @@
+import {
+  ROOT_CONTEXT,
+  context,
+  trace,
+  type Context,
+  type ContextManager,
+} from "@opentelemetry/api";
+import type { ReadableSpan, SpanExporter } from "@opentelemetry/sdk-trace-base";
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
+import {
+  assistantCloudTraceExportOptions,
+  assistantCloudTraceMetadata,
+  createAssistantCloudSpanProcessor,
+  isAssistantCloudSpan,
+  withAssistantCloudTraceMetadata,
+} from "./index";
+
+const TRACE_ID = "0123456789abcdef0123456789abcdef";
+const SPAN_ID = "0123456789abcdef";
+
+class TestContextManager implements ContextManager {
+  private activeContext: Context = ROOT_CONTEXT;
+
+  active(): Context {
+    return this.activeContext;
+  }
+
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+    nextContext: Context,
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    const previousContext = this.activeContext;
+    this.activeContext = nextContext;
+    try {
+      return fn.call(thisArg, ...args);
+    } finally {
+      this.activeContext = previousContext;
+    }
+  }
+
+  bind<T>(_context: Context, target: T): T {
+    return target;
+  }
+
+  enable(): this {
+    return this;
+  }
+
+  disable(): this {
+    this.activeContext = ROOT_CONTEXT;
+    return this;
+  }
+}
+
+function span(
+  name: string,
+  attributes: Record<string, unknown> = {},
+): ReadableSpan {
+  return {
+    name,
+    attributes,
+    spanContext: () => ({
+      traceId: TRACE_ID,
+      spanId: SPAN_ID,
+      traceFlags: 1,
+    }),
+    resource: { asyncAttributesPending: false },
+  } as ReadableSpan;
+}
+
+function recordingExporter() {
+  const spans: ReadableSpan[] = [];
+  const exporter = {
+    export: vi.fn(
+      (
+        nextSpans: ReadableSpan[],
+        resultCallback: Parameters<SpanExporter["export"]>[1],
+      ) => {
+        spans.push(...nextSpans);
+        resultCallback({ code: 0 });
+      },
+    ),
+    shutdown: vi.fn().mockResolvedValue(undefined),
+  } satisfies SpanExporter;
+
+  return { exporter, spans };
+}
+
+beforeAll(() => {
+  context.disable();
+  context.setGlobalContextManager(new TestContextManager());
+});
+
+afterAll(() => {
+  context.disable();
+});
+
+describe("assistantCloudTraceExportOptions", () => {
+  it("uses the default receiver and authorization header", () => {
+    expect(assistantCloudTraceExportOptions({ apiKey: "key" })).toEqual({
+      url: "https://backend.assistant-api.com/v1/traces",
+      headers: { Authorization: "Bearer key" },
+    });
+  });
+
+  it("uses a custom receiver without a trailing slash", () => {
+    expect(
+      assistantCloudTraceExportOptions({
+        apiKey: "key",
+        baseUrl: "https://cloud.example.com",
+      }),
+    ).toEqual({
+      url: "https://cloud.example.com/v1/traces",
+      headers: { Authorization: "Bearer key" },
+    });
+  });
+
+  it("removes trailing slashes and preserves extra headers", () => {
+    expect(
+      assistantCloudTraceExportOptions({
+        apiKey: "key",
+        baseUrl: "https://cloud.example.com/",
+        headers: { "X-Workspace": "workspace_1" },
+      }),
+    ).toEqual({
+      url: "https://cloud.example.com/v1/traces",
+      headers: {
+        Authorization: "Bearer key",
+        "X-Workspace": "workspace_1",
+      },
+    });
+  });
+
+  it("rejects an empty API key", () => {
+    expect(() => assistantCloudTraceExportOptions({ apiKey: "" })).toThrow(
+      "An Assistant Cloud API key is required",
+    );
+  });
+});
+
+describe("isAssistantCloudSpan", () => {
+  it("accepts AI span name prefixes", () => {
+    expect(isAssistantCloudSpan(span("gen_ai.client"))).toBe(true);
+    expect(isAssistantCloudSpan(span("ai.streamText"))).toBe(true);
+    expect(isAssistantCloudSpan(span("llm.generate"))).toBe(true);
+  });
+
+  it("accepts the GenAI operation attribute", () => {
+    expect(
+      isAssistantCloudSpan(
+        span("generate", { "gen_ai.operation.name": "ai.generateText" }),
+      ),
+    ).toBe(true);
+  });
+
+  it("rejects ordinary spans", () => {
+    expect(isAssistantCloudSpan(span("GET /api/chat"))).toBe(false);
+  });
+});
+
+describe("createAssistantCloudSpanProcessor", () => {
+  it("exports matching spans and drops other spans", async () => {
+    const { exporter, spans } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter);
+
+    processor.onEnd(span("gen_ai.client"));
+    processor.onEnd(span("GET /api/chat"));
+    await processor.forceFlush();
+
+    expect(spans.map((candidate) => candidate.name)).toEqual(["gen_ai.client"]);
+  });
+
+  it("uses a custom filter when one is provided", async () => {
+    const { exporter, spans } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter, {
+      filter: (candidate) => candidate.name === "custom",
+    });
+
+    processor.onEnd(span("gen_ai.client"));
+    processor.onEnd(span("custom"));
+    await processor.forceFlush();
+
+    expect(spans.map((candidate) => candidate.name)).toEqual(["custom"]);
+  });
+
+  it("flushes and shuts down the wrapped processor", async () => {
+    const { exporter } = recordingExporter();
+    const processor = createAssistantCloudSpanProcessor(exporter);
+
+    await processor.forceFlush();
+    await processor.shutdown();
+
+    expect(exporter.shutdown).toHaveBeenCalledTimes(1);
+  });
+});
+
+describe("assistantCloudTraceMetadata", () => {
+  it("returns the active trace ID", () => {
+    const result = context.with(
+      trace.setSpanContext(context.active(), {
+        traceId: TRACE_ID,
+        spanId: SPAN_ID,
+        traceFlags: 1,
+      }),
+      assistantCloudTraceMetadata,
+    );
+
+    expect(result).toEqual({ traceId: TRACE_ID });
+  });
+
+  it("returns an empty object without an active span", () => {
+    expect(assistantCloudTraceMetadata()).toEqual({});
+  });
+});
+
+describe("withAssistantCloudTraceMetadata", () => {
+  it("merges the active trace ID into start metadata", () => {
+    const messageMetadata = withAssistantCloudTraceMetadata(({ part }) => ({
+      type: part.type,
+    }));
+
+    const result = context.with(
+      trace.setSpanContext(context.active(), {
+        traceId: TRACE_ID,
+        spanId: SPAN_ID,
+        traceFlags: 1,
+      }),
+      () => messageMetadata({ part: { type: "start" } }),
+    );
+
+    expect(result).toEqual({ type: "start", traceId: TRACE_ID });
+  });
+
+  it("passes through start metadata when no span is active", () => {
+    const value = { started: true };
+    const messageMetadata = withAssistantCloudTraceMetadata(() => value);
+
+    expect(messageMetadata({ part: { type: "start" } })).toBe(value);
+    expect(withAssistantCloudTraceMetadata()({ part: { type: "start" } })).toBe(
+      undefined,
+    );
+  });
+
+  it("passes through metadata for other parts", () => {
+    const value = { usage: { totalTokens: 42 } };
+    const messageMetadata = withAssistantCloudTraceMetadata(() => value);
+
+    expect(messageMetadata({ part: { type: "finish" } })).toBe(value);
+  });
+});

package/src/telemetry/index.ts

@@ -0,0 +1,98 @@
+import { trace } from "@opentelemetry/api";
+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
+import {
+  BatchSpanProcessor,
+  type ReadableSpan,
+  type SpanExporter,
+  type SpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+
+const DEFAULT_BACKEND_BASE_URL = "https://backend.assistant-api.com";
+const AI_SPAN_PREFIXES = ["gen_ai.", "ai.", "llm."];
+
+export type AssistantCloudTraceExportOptions = {
+  apiKey: string;
+  baseUrl?: string;
+  headers?: Record<string, string>;
+};
+
+export type AssistantCloudSpanProcessorOptions = {
+  filter?: (span: ReadableSpan) => boolean;
+};
+
+export function assistantCloudTraceExportOptions({
+  apiKey,
+  baseUrl = DEFAULT_BACKEND_BASE_URL,
+  headers,
+}: AssistantCloudTraceExportOptions): {
+  url: string;
+  headers: Record<string, string>;
+} {
+  if (!apiKey) throw new Error("An Assistant Cloud API key is required");
+
+  return {
+    url: `${trimTrailingSlashes(baseUrl)}/v1/traces`,
+    headers: {
+      Authorization: `Bearer ${apiKey}`,
+      ...headers,
+    },
+  };
+}
+
+function trimTrailingSlashes(url: string): string {
+  let end = url.length;
+  while (end > 0 && url[end - 1] === "/") end--;
+  return url.slice(0, end);
+}
+
+export function createAssistantCloudTraceExporter(
+  options: AssistantCloudTraceExportOptions,
+): SpanExporter {
+  return new OTLPTraceExporter(assistantCloudTraceExportOptions(options));
+}
+
+export function isAssistantCloudSpan(span: ReadableSpan): boolean {
+  return (
+    typeof span.attributes["gen_ai.operation.name"] === "string" ||
+    AI_SPAN_PREFIXES.some((prefix) => span.name.startsWith(prefix))
+  );
+}
+
+export function createAssistantCloudSpanProcessor(
+  exporter: SpanExporter,
+  options: AssistantCloudSpanProcessorOptions = {},
+): SpanProcessor {
+  const processor = new BatchSpanProcessor(exporter);
+  const filter = options.filter ?? isAssistantCloudSpan;
+
+  return {
+    onStart: (span, parentContext) => processor.onStart(span, parentContext),
+    onEnd: (span) => {
+      if (filter(span)) processor.onEnd(span);
+    },
+    forceFlush: () => processor.forceFlush(),
+    shutdown: () => processor.shutdown(),
+  };
+}
+
+export function assistantCloudTraceMetadata(): { traceId?: string } {
+  const spanContext = trace.getActiveSpan()?.spanContext();
+  return spanContext && trace.isSpanContextValid(spanContext)
+    ? { traceId: spanContext.traceId }
+    : {};
+}
+
+export function withAssistantCloudTraceMetadata<
+  Part extends { type: string } = { type: string },
+>(
+  messageMetadata?: (options: {
+    part: Part;
+  }) => Record<string, unknown> | undefined,
+): (options: { part: Part }) => Record<string, unknown> | undefined {
+  return (options) => {
+    const metadata = messageMetadata?.(options);
+    if (options.part.type !== "start") return metadata;
+    const trace = assistantCloudTraceMetadata();
+    return trace.traceId === undefined ? metadata : { ...metadata, ...trace };
+  };
+}