Add trace sampling support (head + tail) for logfire-node and logfire-browser

Implements a two-layer sampling system matching the Python SDK:
- Head sampling: probabilistic at trace creation via ParentBasedSampler
- Tail sampling: callback-based with span buffering via TailSamplingProcessor
- SamplingOptions type, SpanLevel class, checkTraceIdRatio, levelOrDuration factory
- LOGFIRE_TRACE_SAMPLE_RATE env var support in logfire-node
- Example script in examples/node/sampling.ts

Author: petyosi

examples/node/package.json

@@ -7,6 +7,7 @@
   "scripts": {
     "start": "node --experimental-strip-types --disable-warning=ExperimentalWarning index.ts",
     "start-and-throw-error": "TRIGGER_ERROR=true node --experimental-strip-types --disable-warning=ExperimentalWarning index.ts",
+    "sampling": "node --experimental-strip-types --disable-warning=ExperimentalWarning sampling.ts",
     "debug": "node --experimental-strip-types --disable-warning=ExperimentalWarning --inspect-brk index.ts",
     "typecheck": "tsc"
   },

examples/node/sampling.ts

@@ -0,0 +1,64 @@
+import * as logfire from '@pydantic/logfire-node'
+
+// --- Configuration ---
+// This example demonstrates three sampling modes. Uncomment the one you want to try.
+
+// MODE 1: Head sampling only — keep 50% of traces (probabilistic, decided at trace creation)
+// logfire.configure({
+//   serviceName: 'sampling-demo',
+//   console: true,
+//   sendToLogfire: false,
+//   sampling: { head: 0.5 },
+// })
+
+// MODE 2: Tail sampling only — keep traces that have warning+ spans or last > 2 seconds
+// logfire.configure({
+//   serviceName: 'sampling-demo',
+//   console: true,
+//   sendToLogfire: false,
+//   sampling: logfire.levelOrDuration({ levelThreshold: 'warning', durationThreshold: 2.0 }),
+// })
+
+// MODE 3: Tail sampling with a custom callback
+logfire.configure({
+  serviceName: 'sampling-demo',
+  sampling: {
+    tail: (spanInfo) => {
+      // Keep any trace that contains an error-level span
+      if (spanInfo.level.gte('error')) return 1.0
+      // Keep traces running longer than 1.5 seconds
+      if (spanInfo.duration > 1.5) return 1.0
+      // Drop everything else
+      return 0.0
+    },
+  },
+})
+
+console.log('\n--- Trace 1: quick info-level trace (should be DROPPED) ---')
+logfire.span('fast operation', {
+  callback: () => {
+    logfire.info('all good here')
+  },
+})
+
+console.log('\n--- Trace 2: trace with an error (should be KEPT) ---')
+logfire.span('operation with error', {
+  callback: () => {
+    logfire.info('starting work')
+    logfire.error('something broke', { detail: 'disk full' })
+    logfire.info('attempted recovery')
+  },
+})
+
+console.log('\n--- Trace 3: slow trace (should be KEPT after delay) ---')
+await logfire.span('slow operation', {
+  callback: async () => {
+    logfire.info('step 1')
+    await new Promise((resolve) => setTimeout(resolve, 2000))
+    logfire.info('step 2 - this triggers the duration threshold')
+  },
+})
+
+console.log('\n--- Done ---')
+console.log('Check the console output above. You should see spans from traces 2 and 3,')
+console.log('but NOT from trace 1 (it was dropped by tail sampling).')

packages/logfire-api/src/TailSamplingProcessor.ts

@@ -0,0 +1,130 @@
+import { Context, HrTime } from '@opentelemetry/api'
+import { ReadableSpan, Span, SpanProcessor } from '@opentelemetry/sdk-trace-base'
+
+import { checkTraceIdRatio, SpanLevel, type TailSamplingSpanInfo } from './sampling'
+
+interface BufferedSpan {
+  context: Context | null
+  event: 'end' | 'start'
+  span: ReadableSpan | Span
+}
+
+interface TraceBuffer {
+  spans: BufferedSpan[]
+  startTime: HrTime
+}
+
+// Sentinel value indicating a buffer was already flushed
+const FLUSHED = Symbol('flushed')
+
+type TailCallback = (spanInfo: TailSamplingSpanInfo) => number
+
+function hrTimeToSeconds(hrTime: HrTime): number {
+  return hrTime[0] + hrTime[1] / 1e9
+}
+
+export class TailSamplingProcessor implements SpanProcessor {
+  private buffers = new Map<string, TraceBuffer | typeof FLUSHED>()
+  private tail: TailCallback
+  private wrapped: SpanProcessor
+
+  constructor(wrapped: SpanProcessor, tail: TailCallback) {
+    this.wrapped = wrapped
+    this.tail = tail
+  }
+
+  async forceFlush(): Promise<void> {
+    return this.wrapped.forceFlush()
+  }
+
+  onEnd(span: ReadableSpan): void {
+    const traceId = span.spanContext().traceId
+    const entry = this.buffers.get(traceId)
+
+    if (entry === FLUSHED) {
+      this.wrapped.onEnd(span)
+      return
+    }
+
+    if (!entry) {
+      this.wrapped.onEnd(span)
+      return
+    }
+
+    entry.spans.push({ context: null, event: 'end', span })
+
+    const isRoot = !span.parentSpanContext
+    if (isRoot) {
+      // Root span ended — check one last time, then discard if still buffered
+      if (!this.checkSpan(span, null, 'end', entry)) {
+        this.buffers.delete(traceId)
+      }
+      return
+    }
+
+    this.checkSpan(span, null, 'end', entry)
+  }
+
+  onStart(span: Span, parentContext: Context): void {
+    const traceId = span.spanContext().traceId
+    const entry = this.buffers.get(traceId)
+
+    if (entry === FLUSHED) {
+      this.wrapped.onStart(span, parentContext)
+      return
+    }
+
+    const readable = span as unknown as ReadableSpan
+    const isRoot = !readable.parentSpanContext
+    if (isRoot && !entry) {
+      const buffer: TraceBuffer = { spans: [], startTime: readable.startTime }
+      this.buffers.set(traceId, buffer)
+      buffer.spans.push({ context: parentContext, event: 'start', span })
+      this.checkSpan(span as unknown as ReadableSpan, parentContext, 'start', buffer)
+      return
+    }
+
+    if (entry) {
+      entry.spans.push({ context: parentContext, event: 'start', span })
+      this.checkSpan(span as unknown as ReadableSpan, parentContext, 'start', entry)
+      return
+    }
+
+    // No buffer and not root — trace started before this processor was active
+    this.wrapped.onStart(span, parentContext)
+  }
+
+  async shutdown(): Promise<void> {
+    this.buffers.clear()
+    return this.wrapped.shutdown()
+  }
+
+  private checkSpan(span: ReadableSpan, context: Context | null, event: 'end' | 'start', buffer: TraceBuffer): boolean {
+    const duration = hrTimeToSeconds(span.startTime) - hrTimeToSeconds(buffer.startTime)
+    const level = SpanLevel.fromSpan(span)
+
+    const info: TailSamplingSpanInfo = { context, duration, event, level, span }
+    const rate = this.tail(info)
+
+    if (rate >= 1.0 || (rate > 0.0 && checkTraceIdRatio(span.spanContext().traceId, rate))) {
+      this.flushBuffer(span.spanContext().traceId, buffer)
+      return true
+    }
+
+    return false
+  }
+
+  private flushBuffer(traceId: string, buffer: TraceBuffer): void {
+    this.buffers.set(traceId, FLUSHED)
+
+    for (const { context, event, span } of buffer.spans) {
+      if (event === 'start') {
+        // context is always set for 'start' events
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        this.wrapped.onStart(span as Span, context!)
+      } else {
+        this.wrapped.onEnd(span as ReadableSpan)
+      }
+    }
+  }
+}

packages/logfire-api/src/constants.ts

@@ -13,6 +13,7 @@ export const ATTRIBUTES_MESSAGE_KEY = `${LOGFIRE_ATTRIBUTES_NAMESPACE}.msg`
 export const ATTRIBUTES_SCRUBBED_KEY = `${LOGFIRE_ATTRIBUTES_NAMESPACE}.scrubbed`
 /** Key for exception fingerprint used for issue grouping */
 export const ATTRIBUTES_EXCEPTION_FINGERPRINT_KEY = `${LOGFIRE_ATTRIBUTES_NAMESPACE}.exception.fingerprint`
+export const ATTRIBUTES_SAMPLE_RATE_KEY = `${LOGFIRE_ATTRIBUTES_NAMESPACE}.sample_rate`
 export const DEFAULT_OTEL_SCOPE = 'logfire'
 export const JSON_SCHEMA_KEY = 'logfire.json_schema'
 export const JSON_NULL_FIELDS_KEY = 'logfire.null_args'

packages/logfire-api/src/index.ts

@@ -16,13 +16,17 @@ import { computeFingerprint } from './fingerprint'
 import { logfireFormatWithExtras } from './formatter'
 import { logfireApiConfig, serializeAttributes } from './logfireApiConfig'
 import * as logfireApiConfigExports from './logfireApiConfig'
+import * as samplingExports from './sampling'
+import { TailSamplingProcessor } from './TailSamplingProcessor'
 import * as ULIDGeneratorExports from './ULIDGenerator'
 
 export * from './AttributeScrubber'
 export { canonicalizeError, computeFingerprint } from './fingerprint'
 export { configureLogfireApi, logfireApiConfig, resolveBaseUrl, resolveSendToLogfire } from './logfireApiConfig'
 export type { LogfireApiConfig, LogfireApiConfigOptions, ScrubbingOptions } from './logfireApiConfig'
+export * from './sampling'
 export { serializeAttributes } from './serializeAttributes'
+export { TailSamplingProcessor } from './TailSamplingProcessor'
 export * from './ULIDGenerator'
 
 export const Level = {
@@ -215,10 +219,12 @@ export function reportError(message: string, error: Error, extraAttributes: Reco
 const defaultExport = {
   ...AttributeScrubbingExports,
   ...fingerprintExports,
+  ...samplingExports,
   ...ULIDGeneratorExports,
   ...logfireApiConfigExports,
 
   serializeAttributes,
+  TailSamplingProcessor,
   Level,
   startSpan,
   span,