workers-axiom 0.1.0

package/src/protocol/index.ts

@@ -0,0 +1,104 @@
+/**
+ * Wire format shared by the producer (logger) and the tail-worker consumer.
+ *
+ * The producer emits JSON lines via `console.log`; the tail worker parses
+ * them and forwards to Axiom. Both sides import these types so the format
+ * is enforced by the type system rather than coordinated by string matching.
+ */
+
+export const SPAN_EVENT_TYPE = 'span' as const
+
+/** Tag identifying a record carrying fields to merge onto `invocation_summary`. */
+export const SUMMARY_PROPERTIES_TYPE = 'summary_properties' as const
+
+export type SpanKind = 'internal' | 'server' | 'client' | 'producer' | 'consumer'
+
+export type SpanStatus = 'ok' | 'error'
+
+export type AttributeValue = string | number | boolean
+
+export type SpanException = {
+  type?: string
+  message: string
+  stacktrace?: string
+  escaped?: boolean
+}
+
+export interface SpanEvent {
+  type: typeof SPAN_EVENT_TYPE
+  trace_id: string
+  span_id: string
+  parent_span_id?: string
+  name: string
+  kind: SpanKind
+  /** Unix epoch nanoseconds, encoded as string to avoid JS number precision loss. */
+  startTimeUnixNano: string
+  endTimeUnixNano: string
+  status: SpanStatus
+  statusMessage?: string
+  service: string
+  environment?: string
+  attributes?: Record<string, AttributeValue>
+  /**
+   * Sampling verdict, fixed at trace root and propagated to every child span.
+   * The tail worker forwards spans to Axiom's traces store only when true.
+   */
+  sampled: boolean
+  /**
+   * Optional exception info attached when status === 'error'.
+   * Field names match OTel semantic conventions (`exception.type`,
+   * `exception.message`, `exception.stacktrace`, `exception.escaped`) so
+   * the tail worker maps them directly into the OTLP `exception` span event.
+   */
+  exception?: SpanException
+}
+
+export function isSpanEvent(value: unknown): value is SpanEvent {
+  if (typeof value !== 'object' || value === null) return false
+  const v = value as Record<string, unknown>
+  return (
+    v.type === SPAN_EVENT_TYPE &&
+    typeof v.trace_id === 'string' &&
+    typeof v.span_id === 'string' &&
+    typeof v.name === 'string' &&
+    typeof v.startTimeUnixNano === 'string' &&
+    typeof v.endTimeUnixNano === 'string' &&
+    typeof v.service === 'string'
+  )
+}
+
+const TRACEPARENT_RE = /^00-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/
+
+export interface TraceParent {
+  trace_id: string
+  span_id: string
+  sampled: boolean
+}
+
+export function parseTraceparent(header: string | null | undefined): TraceParent | undefined {
+  if (!header) return undefined
+  const match = TRACEPARENT_RE.exec(header.trim().toLowerCase())
+  if (!match) return undefined
+  const [, trace_id, span_id, flags] = match
+  return { trace_id: trace_id!, span_id: span_id!, sampled: (parseInt(flags!, 16) & 1) === 1 }
+}
+
+export function formatTraceparent(trace_id: string, span_id: string, sampled: boolean): string {
+  return `00-${trace_id}-${span_id}-${sampled ? '01' : '00'}`
+}
+
+export function generateTraceId(): string {
+  return randomHex(32)
+}
+
+export function generateSpanId(): string {
+  return randomHex(16)
+}
+
+function randomHex(length: number): string {
+  const bytes = new Uint8Array(length / 2)
+  crypto.getRandomValues(bytes)
+  let out = ''
+  for (const b of bytes) out += b.toString(16).padStart(2, '0')
+  return out
+}

package/src/tail/index.ts

@@ -0,0 +1,237 @@
+import { isSpanEvent, type SpanEvent, SUMMARY_PROPERTIES_TYPE } from '../protocol'
+import { sendSpansToAxiom } from './spans'
+
+export interface AxiomConfig {
+  /** Axiom API token. */
+  axiomToken: string
+  /** Axiom dataset receiving both ingest events and OTLP traces. */
+  axiomDataset: string
+  /**
+   * Axiom ingest endpoint. Defaults to the US edge:
+   * `https://us-east-1.aws.edge.axiom.co/v1/ingest`.
+   * Override with the EU edge or self-hosted Axiom URL as needed.
+   * The dataset name is appended automatically.
+   */
+  ingestBaseUrl?: string
+  /**
+   * Axiom OTLP traces endpoint. Defaults to `https://api.axiom.co/v1/traces`.
+   */
+  tracesEndpoint?: string
+}
+
+export type TailHandler = (
+  events: TraceItem[],
+  env: unknown,
+  ctx: ExecutionContext
+) => Promise<void>
+
+/**
+ * Build a Cloudflare Workers `tail()` handler that forwards logs, metrics,
+ * errors, and spans from any number of producing workers to Axiom.
+ *
+ * Configure each producing worker with:
+ *   "tail_consumers": [{ "service": "<name-of-this-tail-worker>" }]
+ *
+ * @example
+ * ```ts
+ * import { createTailHandler } from 'workers-axiom/tail'
+ *
+ * interface Env { AXIOM_TOKEN: string; AXIOM_DATASET: string }
+ *
+ * export default {
+ *   tail: (events, env: Env, ctx) =>
+ *     createTailHandler({
+ *       axiomToken: env.AXIOM_TOKEN,
+ *       axiomDataset: env.AXIOM_DATASET,
+ *     })(events, env, ctx),
+ * }
+ * ```
+ */
+export function createTailHandler(config: AxiomConfig): TailHandler {
+  const { axiomToken, axiomDataset } = config
+  const ingestBase = config.ingestBaseUrl ?? 'https://us-east-1.aws.edge.axiom.co/v1/ingest'
+
+  return async function tail(events, _env, _ctx) {
+    const axiomEvents: AxiomEvent[] = []
+    const spans: SpanEvent[] = []
+
+    for (const event of events) {
+      const worker = event.scriptName ?? 'unknown'
+      const eventTime = new Date(event.eventTimestamp ?? Date.now()).toISOString()
+
+      // Fields the app explicitly opted into surfacing on `invocation_summary`,
+      // collected from `summary_properties` records emitted via `logger.summary`.
+      // Last-write-wins.
+      const summaryProps: Record<string, unknown> = {}
+      // Last trace_id observed in this worker invocation. A single invocation
+      // can contain multiple traces; we keep the last one because it's the
+      // most specific to the invocation's outbound work.
+      let traceId: string | undefined
+
+      for (const log of event.logs) {
+        const logTime = new Date(log.timestamp).toISOString()
+
+        for (const msg of log.message) {
+          if (typeof msg !== 'string') continue
+
+          let parsed: unknown
+          try {
+            parsed = JSON.parse(msg)
+          } catch {
+            continue
+          }
+
+          if (typeof parsed !== 'object' || parsed === null) continue
+
+          // Span events go to Axiom's /v1/traces endpoint, not the ingest
+          // dataset. Don't forward as a regular log row.
+          if (isSpanEvent(parsed)) {
+            spans.push(parsed)
+            traceId = parsed.trace_id
+            continue
+          }
+
+          const record = parsed as Record<string, unknown>
+
+          if (typeof record.trace_id === 'string') {
+            traceId = record.trace_id
+          }
+
+          if (record.type === SUMMARY_PROPERTIES_TYPE) {
+            for (const [key, value] of Object.entries(record)) {
+              if (key === 'type' || key === 'trace_id') continue
+              summaryProps[key] = value
+            }
+            continue
+          }
+
+          axiomEvents.push({
+            _time: logTime,
+            worker,
+            ...record,
+          })
+        }
+      }
+
+      const requestInfo = extractRequestInfo(event.event)
+
+      axiomEvents.push({
+        _time: eventTime,
+        worker,
+        type: 'invocation_summary',
+        trace_id: traceId,
+        ...summaryProps,
+        ...requestInfo,
+        cpuTime: event.cpuTime,
+        wallTime: event.wallTime,
+        traceEvent: event.event as unknown,
+      })
+    }
+
+    // Run synchronously (not via waitUntil) so failures surface in tail logs
+    // and OTLP/ingest errors aren't silently swallowed by the runtime.
+    await Promise.all([
+      sendToAxiom({ events: axiomEvents, axiomToken, axiomDataset, ingestBase }),
+      sendSpansToAxiom({
+        spans,
+        axiomToken,
+        axiomDataset,
+        ...(config.tracesEndpoint !== undefined && { tracesEndpoint: config.tracesEndpoint }),
+      }),
+    ])
+  }
+}
+
+interface AxiomEvent {
+  _time: string
+  worker: string
+  [key: string]: unknown
+}
+
+interface TraceItemFetchEvent {
+  request?: {
+    method?: string
+    url?: string
+    headers?: Record<string, string>
+    cf?: {
+      country?: string
+      city?: string
+      asn?: number
+    }
+  }
+  response?: {
+    status?: number
+  }
+}
+
+interface ExtractedFields {
+  reqMethod?: string
+  reqUrl?: string
+  reqPath?: string
+  reqHost?: string
+  reqIp?: string
+  reqCountry?: string
+  reqCity?: string
+  reqAsn?: number
+  reqUserAgent?: string
+  reqCfRay?: string
+  resStatus?: number
+}
+
+function extractRequestInfo(event: unknown): ExtractedFields | null {
+  const fetchEvent = event as TraceItemFetchEvent | undefined
+  if (!fetchEvent?.request) return null
+
+  const req = fetchEvent.request
+  const headers = req.headers ?? {}
+  const cf = req.cf ?? {}
+
+  let path: string | undefined
+  try {
+    if (req.url) path = new URL(req.url).pathname
+  } catch {
+    // Invalid URL
+  }
+
+  return {
+    reqMethod: req.method,
+    reqUrl: req.url,
+    reqPath: path,
+    reqHost: headers['host'],
+    reqIp: headers['cf-connecting-ip'],
+    reqCountry: cf.country,
+    reqCity: cf.city,
+    reqAsn: cf.asn,
+    reqUserAgent: headers['user-agent'],
+    reqCfRay: headers['cf-ray'],
+    resStatus: fetchEvent.response?.status,
+  }
+}
+
+async function sendToAxiom(params: {
+  events: AxiomEvent[]
+  axiomToken: string
+  axiomDataset: string
+  ingestBase: string
+}): Promise<void> {
+  const { events, axiomToken, axiomDataset, ingestBase } = params
+  if (events.length === 0) return
+
+  const response = await fetch(`${ingestBase}/${axiomDataset}`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${axiomToken}`,
+    },
+    body: JSON.stringify(events),
+  })
+
+  if (!response.ok) {
+    const text = await response.text().catch(() => '<no body>')
+    throw new Error(
+      `Failed to send logs to Axiom: ${response.status} ${response.statusText} — ${text}`
+    )
+  }
+}
+
+export { sendSpansToAxiom } from './spans'

package/src/tail/spans.ts

@@ -0,0 +1,137 @@
+/**
+ * Span forwarding to Axiom's OTLP `/v1/traces` endpoint.
+ *
+ * Spans arrive as `console.log` JSON events with the {@link SpanEvent} wire format.
+ * The sampling verdict is decided once at the trace root by the producing logger
+ * and stamped on every span via `SpanEvent.sampled`; the tail worker is a pure
+ * router that forwards only `sampled === true` spans. Logs/metrics/errors are
+ * forwarded unconditionally by the caller — sampling gates traces, not log visibility.
+ */
+import type { SpanEvent } from '../protocol'
+
+const SPAN_KIND_TO_OTLP: Record<SpanEvent['kind'], number> = {
+  internal: 1,
+  server: 2,
+  client: 3,
+  producer: 4,
+  consumer: 5,
+}
+
+export interface SendSpansInput {
+  spans: SpanEvent[]
+  axiomToken: string
+  axiomDataset: string
+  /** Full URL of Axiom's OTLP traces endpoint. Defaults to `https://api.axiom.co/v1/traces`. */
+  tracesEndpoint?: string
+}
+
+export async function sendSpansToAxiom(input: SendSpansInput): Promise<void> {
+  const { spans, axiomToken, axiomDataset } = input
+  const endpoint = input.tracesEndpoint ?? 'https://api.axiom.co/v1/traces'
+  const kept = spans.filter((s) => s.sampled)
+  if (kept.length === 0) return
+
+  const body = toOtlpPayload(kept)
+  const response = await fetch(endpoint, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${axiomToken}`,
+      'X-Axiom-Dataset': axiomDataset,
+    },
+    body: JSON.stringify(body),
+  })
+
+  if (!response.ok) {
+    const text = await response.text().catch(() => '<no body>')
+    throw new Error(
+      `Failed to send spans to Axiom: ${response.status} ${response.statusText} — ${text}`
+    )
+  }
+}
+
+interface OtlpResourceSpans {
+  resource: { attributes: OtlpAttribute[] }
+  scopeSpans: { scope: { name: string }; spans: OtlpSpan[] }[]
+}
+
+interface OtlpAttribute {
+  key: string
+  value: { stringValue: string } | { intValue: number } | { boolValue: boolean }
+}
+
+interface OtlpSpan {
+  traceId: string
+  spanId: string
+  parentSpanId?: string
+  name: string
+  kind: number
+  startTimeUnixNano: string
+  endTimeUnixNano: string
+  status: { code: number; message?: string }
+  attributes: OtlpAttribute[]
+  events?: { timeUnixNano: string; name: string; attributes: OtlpAttribute[] }[]
+}
+
+/**
+ * Group spans by `(service, environment)` so each unique pair becomes its own
+ * OTLP `ResourceSpans` entry. Axiom's Traces dashboard groups by `service.name`
+ * from the resource block, so this grouping is what makes per-service filtering work.
+ */
+function toOtlpPayload(spans: SpanEvent[]): { resourceSpans: OtlpResourceSpans[] } {
+  const byResource = new Map<string, SpanEvent[]>()
+  for (const span of spans) {
+    const key = `${span.service} ${span.environment ?? ''}`
+    const list = byResource.get(key)
+    if (list) list.push(span)
+    else byResource.set(key, [span])
+  }
+
+  const resourceSpans: OtlpResourceSpans[] = []
+  for (const group of byResource.values()) {
+    const first = group[0]!
+    const resource = {
+      attributes: [
+        attr('service.name', first.service),
+        ...(first.environment ? [attr('deployment.environment', first.environment)] : []),
+      ],
+    }
+    resourceSpans.push({
+      resource,
+      scopeSpans: [{ scope: { name: first.service }, spans: group.map(toOtlpSpan) }],
+    })
+  }
+  return { resourceSpans }
+}
+
+function toOtlpSpan(span: SpanEvent): OtlpSpan {
+  const attributes = Object.entries(span.attributes ?? {}).map(([key, value]) => attr(key, value))
+  const result: OtlpSpan = {
+    traceId: span.trace_id,
+    spanId: span.span_id,
+    parentSpanId: span.parent_span_id,
+    name: span.name,
+    kind: SPAN_KIND_TO_OTLP[span.kind],
+    startTimeUnixNano: span.startTimeUnixNano,
+    endTimeUnixNano: span.endTimeUnixNano,
+    status: { code: span.status === 'error' ? 2 : 1, message: span.statusMessage },
+    attributes,
+  }
+  if (span.exception) {
+    const eventAttrs: OtlpAttribute[] = []
+    for (const [key, value] of Object.entries(span.exception)) {
+      if (value === undefined) continue
+      eventAttrs.push(attr(`exception.${key}`, value as string | number | boolean))
+    }
+    result.events = [
+      { timeUnixNano: span.endTimeUnixNano, name: 'exception', attributes: eventAttrs },
+    ]
+  }
+  return result
+}
+
+function attr(key: string, value: string | number | boolean): OtlpAttribute {
+  if (typeof value === 'string') return { key, value: { stringValue: value } }
+  if (typeof value === 'boolean') return { key, value: { boolValue: value } }
+  return { key, value: { intValue: value } }
+}