Automatically close the spans started with logfire.span

Include metadata about what was scrubbed

Author: petyosi

.changeset/weak-dryers-scream.md

@@ -0,0 +1,8 @@
+---
+"logfire": minor
+"@pydantic/logfire-api": minor
+"@pydantic/logfire-browser": minor
+"@pydantic/logfire-cf-workers": minor
+---
+
+Auto-close spans opened with logfire.span

README.md

@@ -180,8 +180,6 @@ export default async function Home() {
     level: logfire.Level.Warning,
   }, async (span) => {
     logfire.info("Nested info span");
-    // ending the span is necessary to ensure it is reported
-    span.end();
     return <div>Hello</div>;
   });
 }

packages/logfire-api/src/AttributeScrubber.ts

@@ -2,7 +2,7 @@
 
 export type JsonPath = (number | string)[]
 export interface ScrubbedNote {
-  matchedSubstring: string
+  matched_substring: string
   path: JsonPath
 }
 
@@ -136,7 +136,7 @@ export class LogfireAttributeScrubber implements BaseScrubber {
     }
 
     const matchedSubstring = match[0]
-    notes.push({ matchedSubstring, path })
+    notes.push({ matched_substring: matchedSubstring, path })
     return `[Scrubbed due to '${matchedSubstring}']`
   }
 

packages/logfire-api/src/index.test.ts

@@ -1,5 +1,5 @@
 import { trace } from '@opentelemetry/api'
-import { expect, test, vi } from 'vitest'
+import { beforeEach, describe, expect, test, vi } from 'vitest'
 
 import { ATTRIBUTES_LEVEL_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY } from './constants'
 import { info } from './index'
@@ -25,22 +25,50 @@ vi.mock('@opentelemetry/api', () => {
   }
 })
 
-test('formats the message with the passed attributes', () => {
-  info('aha {i}', { i: 1 })
-  const tracer = trace.getTracer('logfire')
-
-  // eslint-disable-next-line @typescript-eslint/unbound-method
-  expect(tracer.startSpan).toBeCalledWith(
-    'aha 1',
-    {
-      attributes: {
-        [ATTRIBUTES_LEVEL_KEY]: 9,
-        [ATTRIBUTES_MESSAGE_TEMPLATE_KEY]: 'aha {i}',
-        [ATTRIBUTES_SPAN_TYPE_KEY]: 'log',
-        [ATTRIBUTES_TAGS_KEY]: [],
-        i: 1,
+describe('info', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+  test('formats the message with the passed attributes', () => {
+    info('aha {i}', { i: 1 })
+    const tracer = trace.getTracer('logfire')
+
+    // eslint-disable-next-line @typescript-eslint/unbound-method
+    expect(tracer.startSpan).toBeCalledWith(
+      'aha 1',
+      {
+        attributes: {
+          [ATTRIBUTES_LEVEL_KEY]: 9,
+          [ATTRIBUTES_MESSAGE_TEMPLATE_KEY]: 'aha {i}',
+          [ATTRIBUTES_SPAN_TYPE_KEY]: 'log',
+          [ATTRIBUTES_TAGS_KEY]: [],
+          i: 1,
+        },
       },
-    },
-    undefined
-  )
+      undefined
+    )
+  })
+
+  test('adds scrubbing details', () => {
+    info('aha {i}', { i: 1, password: 'hunter' })
+    const tracer = trace.getTracer('logfire')
+
+    // eslint-disable-next-line @typescript-eslint/unbound-method
+    expect(tracer.startSpan).toBeCalledWith(
+      'aha 1',
+      {
+        attributes: {
+          [ATTRIBUTES_LEVEL_KEY]: 9,
+          [ATTRIBUTES_MESSAGE_TEMPLATE_KEY]: 'aha {i}',
+          [ATTRIBUTES_SPAN_TYPE_KEY]: 'log',
+          [ATTRIBUTES_TAGS_KEY]: [],
+          i: 1,
+          'logfire.json_schema': '{"properties":{"logfire.scrubbed":{"type":"array"}},"type":"object"}',
+          'logfire.scrubbed': '[{"matched_substring":"password","path":["password"]}]',
+          password: "[Scrubbed due to 'password']",
+        },
+      },
+      undefined
+    )
+  })
 })

packages/logfire-api/src/index.ts

@@ -48,13 +48,12 @@ export function startSpan(
   attributes: Record<string, unknown> = {},
   { log, tags = [], level = Level.Info }: LogOptions = {}
 ): Span {
-  // TODO: should we also send the extra attributes (2nd arg)?
-  const [formattedMessage, , newTemplate] = logfireFormatWithExtras(msgTemplate, attributes, logfireApiConfig.scrubber)
+  const [formattedMessage, extraAttributes, newTemplate] = logfireFormatWithExtras(msgTemplate, attributes, logfireApiConfig.scrubber)
   const span = logfireApiConfig.tracer.startSpan(
     formattedMessage,
     {
       attributes: {
-        ...serializeAttributes(attributes),
+        ...serializeAttributes({ ...attributes, ...extraAttributes }),
         [ATTRIBUTES_MESSAGE_TEMPLATE_KEY]: newTemplate,
         [ATTRIBUTES_LEVEL_KEY]: level,
         [ATTRIBUTES_TAGS_KEY]: Array.from(new Set(tags).values()),
@@ -67,24 +66,29 @@ export function startSpan(
   return span
 }
 
-export function span<F extends (span: Span) => unknown>(
+export function span<R>(
   msgTemplate: string,
   attributes: Record<string, unknown> = {},
   { tags = [], level = Level.Info }: LogOptions = {},
-  callback: F
+  callback: (span: Span) => R
 ) {
-  const [formattedMessage, , newTemplate] = logfireFormatWithExtras(msgTemplate, attributes, logfireApiConfig.scrubber)
-  return logfireApiConfig.tracer.startActiveSpan<F>(
+  const [formattedMessage, extraAttributes, newTemplate] = logfireFormatWithExtras(msgTemplate, attributes, logfireApiConfig.scrubber)
+
+  return logfireApiConfig.tracer.startActiveSpan(
     formattedMessage,
     {
       attributes: {
-        ...serializeAttributes(attributes),
+        ...serializeAttributes({ ...attributes, ...extraAttributes }),
         [ATTRIBUTES_MESSAGE_TEMPLATE_KEY]: newTemplate,
         [ATTRIBUTES_LEVEL_KEY]: level,
         [ATTRIBUTES_TAGS_KEY]: Array.from(new Set(tags).values()),
       },
     },
-    callback
+    (span: Span) => {
+      const result = callback(span)
+      span.end()
+      return result
+    }
   )
 }
 

packages/logfire-api/src/serializeAttributes.ts

@@ -1,5 +1,5 @@
-import { logfireApiConfig } from '.'
-import { ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY, JSON_NULL_FIELDS_KEY, JSON_SCHEMA_KEY } from './constants'
+import { logfireApiConfig, ScrubbedNote } from '.'
+import { ATTRIBUTES_SCRUBBED_KEY, ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY, JSON_NULL_FIELDS_KEY, JSON_SCHEMA_KEY } from './constants'
 
 export type AttributeValue = boolean | number | string | string[]
 
@@ -20,12 +20,19 @@ type SerializedAttributes = Record<string, AttributeValue>
 export function serializeAttributes(attributes: RawAttributes): SerializedAttributes {
   const scrubber = logfireApiConfig.scrubber
   const alreadyScubbed = ATTRIBUTES_SPAN_TYPE_KEY in attributes
-  const scrubbedAttributes = alreadyScubbed ? attributes : (scrubber.scrubValue([], attributes)[0] as Record<string, unknown>)
-  // if the span is created through the logfire API methods, the attributes have already been scrubbed
+  const [scrubbedAttributes, scrubNotes] = alreadyScubbed ? [attributes, []] : scrubber.scrubValue([], attributes)
 
   const result: SerializedAttributes = {}
   const nullArgs: string[] = []
   const schema: JSONSchema = { properties: {}, type: 'object' }
+
+  if (scrubNotes.length > 0) {
+    if (ATTRIBUTES_SCRUBBED_KEY in scrubbedAttributes) {
+      ;(scrubbedAttributes[ATTRIBUTES_SCRUBBED_KEY] as ScrubbedNote[]).push(...scrubNotes)
+    } else {
+      scrubbedAttributes[ATTRIBUTES_SCRUBBED_KEY] = scrubNotes
+    }
+  }
   for (const [key, value] of Object.entries(scrubbedAttributes)) {
     // we don't want to serialize the tags
     if (key === ATTRIBUTES_TAGS_KEY) {