Add setTag to emitter interface to set arbitrary key/value

Here, bring in a `setTag` function that lets X-ray users set an
arbitrary key/value on a request span. These pairs are JSON encoded at
the end of a request and emitted into the span. This is one that we've
already supported for a while in Go and on the server, so we're just
bringing equivalent functionality to this emitter so we can get access
to it from the Stainless stack.

Author: brandur

src/core/attributes.ts

@@ -10,7 +10,7 @@ import {
   ATTR_URL_PATH,
   ATTR_USER_ID,
 } from '@opentelemetry/semantic-conventions/incubating';
-import type { CapturedBody } from './types';
+import type { AttributeValue, CapturedBody } from './types';
 import {
   AttributeKeyRequestBody,
   AttributeKeyRequestBodyEncoding,
@@ -19,6 +19,7 @@ import {
   AttributeKeyResponseBody,
   AttributeKeyResponseBodyEncoding,
   AttributeKeyResponseBodyTruncated,
+  AttributeKeyTags,
   AttributeKeyTenantID,
 } from './attrkey';
 
@@ -291,3 +292,7 @@ export function setTenantIdAttribute(span: Span, tenantId: string): void {
 export function setRequestIdAttribute(span: Span, requestId: string): void {
   span.setAttribute(AttributeKeyRequestID, requestId);
 }
+
+export function setTagsAttribute(span: Span, tags: Record<string, AttributeValue>): void {
+  span.setAttribute(AttributeKeyTags, JSON.stringify(tags));
+}

src/core/attrkey.ts

@@ -8,6 +8,7 @@ export const AttributeKeyResponseBody = 'http.response.body';
 export const AttributeKeyResponseBodyEncoding = 'http.response.body.encoding';
 export const AttributeKeyResponseBodyTruncated = 'http.response.body.truncated';
 export const AttributeKeySpanDrop = 'stainlessxray.internal.drop';
+export const AttributeKeyTags = 'stainlessxray.internal.tags';
 export const AttributeKeyTenantID = 'stainlessxray.tenant.id';
 export const AttributeKeyXrayCaptureErrorCode = 'xray.capture_error_code';
 export const AttributeKeyXrayErrorThrown = 'xray.error_thrown';

src/core/emitter.test.ts

@@ -331,3 +331,99 @@ test('setUserId logs when span attributes fail', () => {
   assert.equal(errors[0]?.msg, 'xray: setUserId failed');
   assert.equal(errors[0]?.fields?.['error'], 'setAttribute failed');
 });
+
+test('setTag records tags in log and span', async () => {
+  const spans: ReadableSpan[] = [];
+  const xray = createEmitter(
+    {
+      serviceName: 'test',
+      endpointUrl: 'https://collector',
+      exporter: { spanProcessor: 'simple' },
+    },
+    createRecordingExporter(spans),
+  );
+
+  const ctx = xray.startRequest({
+    method: 'GET',
+    url: 'https://example.test/tags',
+    headers: {},
+    startTimeMs: 0,
+  });
+  ctx.setTag('environment', 'staging');
+  ctx.setTag('region', 'us-east-1');
+  ctx.setTag('retries', 3);
+  ctx.setTag('verbose', true);
+
+  const log = xray.endRequest(ctx, {
+    statusCode: 200,
+    headers: {},
+    endTimeMs: 5,
+  });
+
+  await xray.flush();
+  assert.deepEqual(log.tags, { environment: 'staging', region: 'us-east-1', retries: 3, verbose: true });
+  const tagsAttr = spans[0]?.attributes['stainlessxray.internal.tags'];
+  assert.equal(typeof tagsAttr, 'string');
+  assert.deepEqual(JSON.parse(tagsAttr as string), {
+    environment: 'staging',
+    region: 'us-east-1',
+    retries: 3,
+    verbose: true,
+  });
+});
+
+test('setTag overwrites previous value for same key', () => {
+  const xray = createEmitter(
+    {
+      serviceName: 'test',
+      endpointUrl: 'https://collector',
+    },
+    createNoopExporter(),
+  );
+
+  const ctx = xray.startRequest({
+    method: 'GET',
+    url: 'https://example.test/tags-overwrite',
+    headers: {},
+    startTimeMs: 0,
+  });
+  ctx.setTag('env', 'dev');
+  ctx.setTag('env', 'prod');
+
+  const log = xray.endRequest(ctx, {
+    statusCode: 200,
+    headers: {},
+    endTimeMs: 5,
+  });
+
+  assert.deepEqual(log.tags, { env: 'prod' });
+});
+
+test('tags omitted from log and span when empty', async () => {
+  const spans: ReadableSpan[] = [];
+  const xray = createEmitter(
+    {
+      serviceName: 'test',
+      endpointUrl: 'https://collector',
+      exporter: { spanProcessor: 'simple' },
+    },
+    createRecordingExporter(spans),
+  );
+
+  const ctx = xray.startRequest({
+    method: 'GET',
+    url: 'https://example.test/no-tags',
+    headers: {},
+    startTimeMs: 0,
+  });
+
+  const log = xray.endRequest(ctx, {
+    statusCode: 200,
+    headers: {},
+    endTimeMs: 5,
+  });
+
+  await xray.flush();
+  assert.equal(log.tags, undefined);
+  assert.equal('stainlessxray.internal.tags' in (spans[0]?.attributes ?? {}), false);
+});

src/core/emitter.ts

@@ -14,6 +14,7 @@ import {
   setResponseBodySizeAttribute,
   setResponseStatusAttribute,
   setRouteAttribute,
+  setTagsAttribute,
   setTenantIdAttribute,
   setUserIdAttribute,
 } from './attributes';
@@ -136,6 +137,13 @@ function startRequest(
       }
       state.sessionId = id;
     },
+    setTag: (key, value) => {
+      const state = getContextState(context);
+      if (!state) {
+        return;
+      }
+      state.tags[key] = value;
+    },
     setAttribute: (key, value) => {
       const state = getContextState(context);
       if (!state) {
@@ -187,6 +195,7 @@ function startRequest(
     context,
     attributes: {},
     events: [],
+    tags: {},
   };
 
   bindContext(context, state);
@@ -250,6 +259,7 @@ function endRequest(
     sessionId: state.sessionId ?? undefined,
     error: buildError(err ?? state.error),
     attributes: Object.keys(state.attributes).length > 0 ? { ...state.attributes } : undefined,
+    tags: Object.keys(state.tags).length > 0 ? { ...state.tags } : undefined,
     timestamp: new Date(endTimeMs).toISOString(),
   };
 
@@ -304,6 +314,9 @@ function endRequest(
       if (state.userId) {
         setUserIdAttribute(span, state.userId);
       }
+      if (Object.keys(state.tags).length > 0) {
+        setTagsAttribute(span, state.tags);
+      }
       if (err ?? state.error) {
         spanStatusFromError(span, err ?? state.error);
       }

src/core/state.ts

@@ -9,6 +9,7 @@ export type RequestState = {
   context: XrayContext;
   attributes: Record<string, AttributeValue>;
   events: Array<{ name: string; attributes?: Record<string, AttributeValue> }>;
+  tags: Record<string, AttributeValue>;
   tenantId?: string;
   userId?: string;
   sessionId?: string;

src/core/types.ts

@@ -119,6 +119,10 @@ export interface RequestLog {
    * Custom attributes added via `setAttribute`.
    */
   attributes?: Record<string, AttributeValue>;
+  /**
+   * Custom tags added via `setTag`.
+   */
+  tags?: Record<string, AttributeValue>;
   /**
    * ISO timestamp for when request processing completed.
    */
@@ -156,6 +160,10 @@ export interface XrayContext {
    * Set a session identifier on the request log.
    */
   setSessionId(id: string): void;
+  /**
+   * Set an arbitrary user-defined key/value tag on the request.
+   */
+  setTag(key: string, value: AttributeValue): void;
   /**
    * Add or overwrite a custom attribute on the request span/log.
    */