@mutagent/sdk 0.2.13

package/src/tracing/span-manager.ts

@@ -0,0 +1,210 @@
+/**
+ * MutagenT SDK — Span Manager
+ *
+ * Manages span lifecycle, context propagation via AsyncLocalStorage,
+ * and parent-child linking for traces.
+ *
+ * NOT generated by Speakeasy — manually maintained.
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+import { randomUUID } from "node:crypto";
+import type {
+  FinishedSpan,
+  MutagentSpan,
+  SerializedSpanEvent,
+  SpanEndOptions,
+  SpanOptions,
+} from "./types.js";
+
+// ---------------------------------------------------------------------------
+// Context — stores current span in async context
+// ---------------------------------------------------------------------------
+
+interface SpanContext {
+  span: MutagentSpan;
+  traceId: string;
+}
+
+const asyncLocalStorage = new AsyncLocalStorage<SpanContext>();
+
+// ---------------------------------------------------------------------------
+// SpanManager
+// ---------------------------------------------------------------------------
+
+export class SpanManager {
+  private readonly activeSpans = new Map<string, MutagentSpan>();
+  private onSpanEnd: ((span: FinishedSpan) => void) | undefined;
+  private debugEnabled = false;
+
+  setOnSpanEnd(callback: (span: FinishedSpan) => void): void {
+    this.onSpanEnd = callback;
+  }
+
+  setDebug(enabled: boolean): void {
+    this.debugEnabled = enabled;
+  }
+
+  /**
+   * Start a new span. If no trace is active, a new trace ID is generated.
+   * The span is automatically linked as a child of the current span if one
+   * exists in the async context.
+   *
+   * Returns the created span. To execute code within the span's context,
+   * use `runInContext()`.
+   */
+  startSpan(options: SpanOptions): MutagentSpan {
+    const currentContext = asyncLocalStorage.getStore();
+
+    // Determine trace ID: reuse active trace or create new
+    const traceId = currentContext?.traceId ?? randomUUID();
+
+    // Determine parent: explicit option > current context > none
+    const parentSpanId =
+      options.parentSpanId ?? currentContext?.span.id ?? undefined;
+
+    const spanId = randomUUID();
+    const span: MutagentSpan = {
+      id: spanId,
+      traceId,
+      parentSpanId,
+      kind: options.kind,
+      name: options.name ?? options.kind,
+      startedAt: new Date(),
+      input: options.input,
+      output: undefined,
+      status: "unset",
+      statusMessage: undefined,
+      metrics: {},
+      attributes: { ...options.attributes },
+      events: [],
+    };
+
+    this.activeSpans.set(spanId, span);
+
+    if (this.debugEnabled) {
+      console.debug(
+        `[mutagent:trace] START span=${spanId} trace=${traceId} kind=${span.kind} name=${span.name}`,
+      );
+    }
+
+    return span;
+  }
+
+  /**
+   * Run a function within the async context of a span.
+   * This enables getCurrentSpan() / getCurrentTraceId() for nested calls.
+   */
+  runInContext<T>(span: MutagentSpan, fn: () => T): T {
+    const context: SpanContext = { span, traceId: span.traceId };
+    return asyncLocalStorage.run(context, fn);
+  }
+
+  /**
+   * End a span: compute duration, build FinishedSpan, remove from active set,
+   * and pass to the BatchCollector via the onSpanEnd callback.
+   */
+  endSpan(span: MutagentSpan, endOptions?: SpanEndOptions): void {
+    const endTime = new Date();
+    const durationMs = endTime.getTime() - span.startedAt.getTime();
+
+    // Apply end options to span
+    if (endOptions?.status !== undefined) {
+      span.status = endOptions.status;
+    }
+    if (endOptions?.statusMessage !== undefined) {
+      span.statusMessage = endOptions.statusMessage;
+    }
+    if (endOptions?.output !== undefined) {
+      span.output = endOptions.output;
+    }
+    if (endOptions?.metrics !== undefined) {
+      Object.assign(span.metrics, endOptions.metrics);
+    }
+    if (endOptions?.attributes !== undefined) {
+      Object.assign(span.attributes, endOptions.attributes);
+    }
+    if (endOptions?.events !== undefined) {
+      span.events.push(...endOptions.events);
+    }
+
+    // Build serialized events
+    const serializedEvents: SerializedSpanEvent[] = span.events.map((e) => ({
+      name: e.name,
+      timestamp: (e.timestamp ?? new Date()).toISOString(),
+      ...(e.attributes !== undefined ? { attributes: e.attributes } : {}),
+    }));
+
+    // Build FinishedSpan
+    const finished: FinishedSpan = {
+      spanId: span.id,
+      traceId: span.traceId,
+      parentSpanId: span.parentSpanId,
+      name: span.name,
+      kind: span.kind,
+      input: span.input,
+      output: span.output,
+      startTime: span.startedAt.toISOString(),
+      endTime: endTime.toISOString(),
+      durationMs,
+      model: span.metrics.model,
+      provider: span.metrics.provider,
+      inputTokens: span.metrics.inputTokens,
+      outputTokens: span.metrics.outputTokens,
+      totalTokens: span.metrics.totalTokens,
+      costUsd: span.metrics.costUsd,
+      status: span.status,
+      statusMessage: span.statusMessage,
+      attributes: span.attributes,
+      events: serializedEvents,
+      metadata: undefined,
+    };
+
+    // Remove from active set
+    this.activeSpans.delete(span.id);
+
+    if (this.debugEnabled) {
+      console.debug(
+        `[mutagent:trace] END   span=${span.id} trace=${span.traceId} status=${finished.status} duration=${durationMs}ms`,
+      );
+    }
+
+    // Deliver to collector
+    this.onSpanEnd?.(finished);
+  }
+
+  /**
+   * Get the currently active span from async context.
+   */
+  getCurrentSpan(): MutagentSpan | undefined {
+    return asyncLocalStorage.getStore()?.span;
+  }
+
+  /**
+   * Get the currently active trace ID from async context.
+   */
+  getCurrentTraceId(): string | undefined {
+    return asyncLocalStorage.getStore()?.traceId;
+  }
+
+  /**
+   * Get an active span by ID.
+   */
+  getActiveSpan(spanId: string): MutagentSpan | undefined {
+    return this.activeSpans.get(spanId);
+  }
+
+  /**
+   * Get count of active spans (useful for tests and debugging).
+   */
+  getActiveSpanCount(): number {
+    return this.activeSpans.size;
+  }
+
+  /**
+   * Clear all active spans (used during shutdown/reset).
+   */
+  clear(): void {
+    this.activeSpans.clear();
+  }
+}

package/src/tracing/types.ts

@@ -0,0 +1,256 @@
+/**
+ * MutagenT SDK — Tracing Type Definitions
+ *
+ * These types define the tracing primitives for the MutagenT SDK.
+ * They align with the server-side trace schema and OTel Gen AI conventions.
+ *
+ * NOT generated by Speakeasy — manually maintained.
+ */
+
+// ---------------------------------------------------------------------------
+// SpanKind — Taxonomy aligned with OTel Gen AI semantic conventions (v1.37+)
+// ---------------------------------------------------------------------------
+
+export type SpanKind =
+  // Generation
+  | "llm.chat"
+  | "llm.completion"
+  | "llm.embedding"
+  // Orchestration
+  | "chain"
+  | "agent"
+  | "graph"
+  | "node"
+  | "edge"
+  | "workflow"
+  | "middleware"
+  // Tools
+  | "tool"
+  // Retrieval
+  | "retrieval"
+  | "rerank"
+  // Safety
+  | "guardrail"
+  // Other
+  | "custom";
+
+// ---------------------------------------------------------------------------
+// SpanIO — Structured Input/Output for spans (stored as JSONB on server)
+// ---------------------------------------------------------------------------
+
+export interface SpanIO {
+  /** Chat messages (most common for LLM spans) */
+  messages?: Array<{
+    role: "system" | "user" | "assistant" | "tool";
+    content: string;
+    name?: string;
+    toolCallId?: string;
+  }>;
+
+  /** Tool calls (for LLM spans that request tool use) */
+  toolCalls?: Array<{
+    id: string;
+    name: string;
+    arguments: Record<string, unknown>;
+    result?: unknown;
+  }>;
+
+  /** Retrieved documents (for retrieval spans) */
+  documents?: Array<{
+    content: string;
+    metadata?: Record<string, unknown>;
+    score?: number;
+  }>;
+
+  /** Raw text (for text completions, generic spans) */
+  text?: string;
+
+  /** Framework-specific raw data (escape hatch) */
+  raw?: unknown;
+}
+
+// ---------------------------------------------------------------------------
+// Span Status
+// ---------------------------------------------------------------------------
+
+export type SpanStatus = "ok" | "error" | "unset";
+
+// ---------------------------------------------------------------------------
+// Span Lifecycle Events
+// ---------------------------------------------------------------------------
+
+export interface SpanEvent {
+  name: string;
+  timestamp?: Date;
+  attributes?: Record<string, unknown>;
+}
+
+// ---------------------------------------------------------------------------
+// Span Metrics — Typed columns for fast queries/aggregation on the server
+// ---------------------------------------------------------------------------
+
+export interface SpanMetrics {
+  model?: string;
+  provider?: string;
+  inputTokens?: number;
+  outputTokens?: number;
+  totalTokens?: number;
+  costUsd?: number;
+  latencyMs?: number;
+  ttftMs?: number;
+}
+
+// ---------------------------------------------------------------------------
+// Span Options — Passed to startSpan() / withTrace() / @trace()
+// ---------------------------------------------------------------------------
+
+export interface SpanOptions {
+  kind: SpanKind;
+  name?: string;
+  input?: SpanIO;
+  attributes?: Record<string, unknown>;
+  parentSpanId?: string;
+}
+
+// ---------------------------------------------------------------------------
+// SpanEndOptions — Passed to endSpan() on span completion
+// ---------------------------------------------------------------------------
+
+export interface SpanEndOptions {
+  status?: SpanStatus;
+  statusMessage?: string;
+  output?: SpanIO;
+  metrics?: SpanMetrics;
+  attributes?: Record<string, unknown>;
+  events?: SpanEvent[];
+}
+
+// ---------------------------------------------------------------------------
+// SpanHandle — Handle passed to withTrace() callback for manual enrichment
+// ---------------------------------------------------------------------------
+
+export interface SpanHandle {
+  readonly spanId: string;
+  readonly traceId: string;
+  setAttributes(attrs: Record<string, unknown>): void;
+  addEvent(name: string, attributes?: Record<string, unknown>): void;
+  setOutput(output: SpanIO): void;
+  setInput(input: SpanIO): void;
+  setMetrics(metrics: SpanMetrics): void;
+}
+
+// ---------------------------------------------------------------------------
+// MutagentSpan — Internal span representation (mutable during lifecycle)
+// ---------------------------------------------------------------------------
+
+export interface MutagentSpan {
+  readonly id: string;
+  readonly traceId: string;
+  readonly parentSpanId: string | undefined;
+  readonly kind: SpanKind;
+  readonly name: string;
+  readonly startedAt: Date;
+  input: SpanIO | undefined;
+  output: SpanIO | undefined;
+  status: SpanStatus;
+  statusMessage: string | undefined;
+  metrics: SpanMetrics;
+  attributes: Record<string, unknown>;
+  events: SpanEvent[];
+}
+
+// ---------------------------------------------------------------------------
+// FinishedSpan — Immutable, ready for serialization and transport
+// ---------------------------------------------------------------------------
+
+export interface FinishedSpan {
+  readonly spanId: string;
+  readonly traceId: string;
+  readonly parentSpanId: string | undefined;
+  readonly name: string;
+  readonly kind: SpanKind;
+  readonly input: SpanIO | undefined;
+  readonly output: SpanIO | undefined;
+  readonly startTime: string; // ISO 8601
+  readonly endTime: string; // ISO 8601
+  readonly durationMs: number;
+  readonly model: string | undefined;
+  readonly provider: string | undefined;
+  readonly inputTokens: number | undefined;
+  readonly outputTokens: number | undefined;
+  readonly totalTokens: number | undefined;
+  readonly costUsd: number | undefined;
+  readonly status: SpanStatus;
+  readonly statusMessage: string | undefined;
+  readonly attributes: Record<string, unknown>;
+  readonly events: SerializedSpanEvent[];
+  readonly metadata: Record<string, unknown> | undefined;
+}
+
+export interface SerializedSpanEvent {
+  readonly name: string;
+  readonly timestamp: string; // ISO 8601
+  readonly attributes?: Record<string, unknown>;
+}
+
+// ---------------------------------------------------------------------------
+// TracingConfig — Configuration for initTracing()
+// ---------------------------------------------------------------------------
+
+export interface TracingConfig {
+  /** API key for authentication. Required. */
+  apiKey: string;
+  /** API endpoint URL. Defaults to 'https://api.mutagent.dev'. */
+  endpoint?: string;
+  /** Environment name (e.g., 'production', 'staging', 'development'). */
+  environment?: string;
+  /** Number of spans to buffer before flushing. Default: 10. */
+  batchSize?: number;
+  /** Flush interval in milliseconds. Default: 5000. */
+  flushIntervalMs?: number;
+  /** Enable debug logging of span start/end to console. Default: false. */
+  debug?: boolean;
+  /** Source identifier for traces. Default: 'sdk'. */
+  source?: string;
+}
+
+// ---------------------------------------------------------------------------
+// TracePayload — HTTP payload sent to the server
+// ---------------------------------------------------------------------------
+
+export interface TracePayload {
+  traces: Array<{
+    traceId: string;
+    name?: string;
+    sessionId?: string;
+    source: string;
+    environment?: string;
+    spans: Array<{
+      spanId: string;
+      parentSpanId?: string;
+      name: string;
+      kind: SpanKind;
+      input?: SpanIO;
+      output?: SpanIO;
+      startTime: string;
+      endTime: string;
+      durationMs: number;
+      model?: string;
+      provider?: string;
+      inputTokens?: number;
+      outputTokens?: number;
+      totalTokens?: number;
+      costUsd?: number;
+      status: SpanStatus;
+      statusMessage?: string;
+      attributes?: Record<string, unknown>;
+      events?: SerializedSpanEvent[];
+      metadata?: Record<string, unknown>;
+    }>;
+  }>;
+}
+
+export interface TraceHTTPResponse {
+  accepted: number;
+  rejected: number;
+}

package/src/tracing/wrappers.ts

@@ -0,0 +1,174 @@
+/**
+ * MutagenT SDK -- withTrace Function Wrapper
+ *
+ * Wraps a function with a tracing span, providing a SpanHandle
+ * for manual enrichment (attributes, events, output, metrics).
+ *
+ * NOT generated by Speakeasy -- manually maintained.
+ */
+
+import type {
+  SpanKind,
+  SpanEvent,
+  SpanHandle,
+  SpanIO,
+  SpanMetrics,
+  MutagentSpan,
+} from "./types.js";
+import { _getSpanManager } from "./init.js";
+
+// ---------------------------------------------------------------------------
+// WithTrace Options
+// ---------------------------------------------------------------------------
+
+export interface WithTraceOptions {
+  kind: SpanKind;
+  name?: string;
+  input?: SpanIO;
+  attributes?: Record<string, unknown>;
+}
+
+// ---------------------------------------------------------------------------
+// SpanHandle implementation
+// ---------------------------------------------------------------------------
+
+class SpanHandleImpl implements SpanHandle {
+  readonly spanId: string;
+  readonly traceId: string;
+  private readonly span: MutagentSpan;
+
+  constructor(span: MutagentSpan) {
+    this.spanId = span.id;
+    this.traceId = span.traceId;
+    this.span = span;
+  }
+
+  setAttributes(attrs: Record<string, unknown>): void {
+    Object.assign(this.span.attributes, attrs);
+  }
+
+  addEvent(name: string, attributes?: Record<string, unknown>): void {
+    const event: SpanEvent = {
+      name,
+      timestamp: new Date(),
+    };
+    if (attributes !== undefined) {
+      event.attributes = attributes;
+    }
+    this.span.events.push(event);
+  }
+
+  setOutput(output: SpanIO): void {
+    this.span.output = output;
+  }
+
+  setInput(input: SpanIO): void {
+    this.span.input = input;
+  }
+
+  setMetrics(metrics: SpanMetrics): void {
+    Object.assign(this.span.metrics, metrics);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// withTrace()
+// ---------------------------------------------------------------------------
+
+/**
+ * Wrap a function with a tracing span.
+ *
+ * The callback receives a SpanHandle for manual enrichment:
+ * - `span.setAttributes(attrs)` -- add custom attributes
+ * - `span.addEvent(name, attrs?)` -- record a lifecycle event
+ * - `span.setOutput(output)` -- set structured output
+ * - `span.setInput(input)` -- set/override structured input
+ * - `span.setMetrics(metrics)` -- set token counts, cost, etc.
+ *
+ * Automatically sets:
+ * - `status: 'ok'` on success
+ * - `status: 'error'` on throw (re-throws the exception)
+ *
+ * No-op when tracing is not initialized (runs fn directly).
+ *
+ * Works with both sync and async functions.
+ *
+ * Usage:
+ * ```ts
+ * const result = await withTrace({ kind: 'chain', name: 'rag' }, async (span) => {
+ *   span.setAttributes({ 'gen_ai.model': 'gpt-4' });
+ *   const answer = await generate(query);
+ *   span.setOutput({ text: answer });
+ *   return answer;
+ * });
+ * ```
+ */
+export function withTrace<T>(
+  options: WithTraceOptions,
+  fn: (span: SpanHandle) => T,
+): T {
+  const manager = _getSpanManager();
+
+  // No-op when tracing is not initialized
+  if (!manager) {
+    // Create a dummy handle so the function signature is satisfied
+    const dummyHandle: SpanHandle = {
+      spanId: "",
+      traceId: "",
+      setAttributes() {},
+      addEvent() {},
+      setOutput() {},
+      setInput() {},
+      setMetrics() {},
+    };
+    return fn(dummyHandle);
+  }
+
+  const spanOptions: import("./types.js").SpanOptions = {
+    kind: options.kind,
+  };
+  if (options.name !== undefined) spanOptions.name = options.name;
+  if (options.input !== undefined) spanOptions.input = options.input;
+  if (options.attributes) spanOptions.attributes = { ...options.attributes };
+
+  const span = manager.startSpan(spanOptions);
+
+  const handle = new SpanHandleImpl(span);
+
+  return manager.runInContext(span, () => {
+    try {
+      const result = fn(handle);
+
+      // Handle async
+      if (result instanceof Promise) {
+        return result.then(
+          (value: unknown) => {
+            manager.endSpan(span, { status: "ok" });
+            return value;
+          },
+          (error: unknown) => {
+            const message =
+              error instanceof Error ? error.message : String(error);
+            manager.endSpan(span, {
+              status: "error",
+              statusMessage: message,
+            });
+            throw error;
+          },
+        ) as T;
+      }
+
+      // Sync
+      manager.endSpan(span, { status: "ok" });
+      return result;
+    } catch (error) {
+      const message =
+        error instanceof Error ? error.message : String(error);
+      manager.endSpan(span, {
+        status: "error",
+        statusMessage: message,
+      });
+      throw error;
+    }
+  });
+}

package/src/types/async.ts

@@ -0,0 +1,68 @@
+/*
+ * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.
+ */
+
+export type APICall =
+  | {
+      status: "complete";
+      request: Request;
+      response: Response;
+    }
+  | {
+      status: "request-error";
+      request: Request;
+      response?: undefined;
+    }
+  | {
+      status: "invalid";
+      request?: undefined;
+      response?: undefined;
+    };
+
+export class APIPromise<T> implements Promise<T> {
+  readonly #promise: Promise<[T, APICall]>;
+  readonly #unwrapped: Promise<T>;
+
+  readonly [Symbol.toStringTag] = "APIPromise";
+
+  constructor(p: [T, APICall] | Promise<[T, APICall]>) {
+    this.#promise = p instanceof Promise ? p : Promise.resolve(p);
+    this.#unwrapped =
+      p instanceof Promise
+        ? this.#promise.then(([value]) => value)
+        : Promise.resolve(p[0]);
+  }
+
+  then<TResult1 = T, TResult2 = never>(
+    onfulfilled?:
+      | ((value: T) => TResult1 | PromiseLike<TResult1>)
+      | null
+      | undefined,
+    onrejected?:
+      | ((reason: any) => TResult2 | PromiseLike<TResult2>)
+      | null
+      | undefined,
+  ): Promise<TResult1 | TResult2> {
+    return this.#promise.then(
+      onfulfilled ? ([value]) => onfulfilled(value) : void 0,
+      onrejected,
+    );
+  }
+
+  catch<TResult = never>(
+    onrejected?:
+      | ((reason: any) => TResult | PromiseLike<TResult>)
+      | null
+      | undefined,
+  ): Promise<T | TResult> {
+    return this.#unwrapped.catch(onrejected);
+  }
+
+  finally(onfinally?: (() => void) | null | undefined): Promise<T> {
+    return this.#unwrapped.finally(onfinally);
+  }
+
+  $inspect(): Promise<[T, APICall]> {
+    return this.#promise;
+  }
+}