@braintrust/temporal 0.1.0

package/src/sinks.ts

@@ -0,0 +1,76 @@
+import type { WorkflowInfo, Sinks } from "@temporalio/workflow";
+import type { InjectedSinks } from "@temporalio/worker";
+import * as braintrust from "braintrust";
+
+// Sink interface (used in workflow code via proxySinks)
+// NOTE: WorkflowInfo is NOT included here - it's automatically injected by the runtime
+export interface BraintrustSinks extends Sinks {
+  braintrust: {
+    workflowStarted(parentContext?: string, workflowSpanId?: string): void;
+    workflowCompleted(error?: string): void;
+  };
+}
+
+// Active workflow spans tracked by run ID
+const workflowSpans = new Map<string, braintrust.Span>();
+// Workflow span exports tracked by run ID (as promises for async export)
+const workflowSpanExports = new Map<string, Promise<string>>();
+
+/**
+ * Get the exported span context for a workflow by run ID.
+ * Activities on the same worker can use this to parent to the workflow span.
+ */
+export function getWorkflowSpanExport(
+  runId: string,
+): Promise<string> | undefined {
+  return workflowSpanExports.get(runId);
+}
+
+/**
+ * Create the Braintrust sinks for workflow span management.
+ * These sinks are called from the workflow isolate via proxySinks.
+ */
+export function createBraintrustSinks(): InjectedSinks<BraintrustSinks> {
+  return {
+    braintrust: {
+      workflowStarted: {
+        fn: (
+          info: WorkflowInfo,
+          parentContext?: string,
+          workflowSpanId?: string,
+        ) => {
+          const span = braintrust.startSpan({
+            name: `temporal.workflow.${info.workflowType}`,
+            spanAttributes: { type: "task" },
+            parent: parentContext,
+            spanId: workflowSpanId,
+            event: {
+              metadata: {
+                "temporal.workflow_type": info.workflowType,
+                "temporal.workflow_id": info.workflowId,
+                "temporal.run_id": info.runId,
+              },
+            },
+          });
+          workflowSpans.set(info.runId, span);
+          workflowSpanExports.set(info.runId, span.export());
+        },
+        callDuringReplay: false,
+      },
+      workflowCompleted: {
+        fn: (info: WorkflowInfo, error?: string) => {
+          const span = workflowSpans.get(info.runId);
+          if (span) {
+            if (error) {
+              span.log({ error });
+            }
+            span.end();
+            workflowSpans.delete(info.runId);
+            workflowSpanExports.delete(info.runId);
+          }
+        },
+        callDuringReplay: false,
+      },
+    },
+  };
+}

package/src/temporal.test.ts

@@ -0,0 +1,243 @@
+import { expect, test, describe } from "vitest";
+import {
+  serializeHeaderValue,
+  deserializeHeaderValue,
+  BRAINTRUST_SPAN_HEADER,
+  BRAINTRUST_WORKFLOW_SPAN_HEADER,
+  BRAINTRUST_WORKFLOW_SPAN_ID_HEADER,
+} from "./utils";
+import { SpanComponentsV3, SpanObjectTypeV3 } from "braintrust/util";
+import {
+  BraintrustTemporalPlugin,
+  createBraintrustTemporalPlugin,
+} from "./plugin";
+
+describe("temporal header utilities", () => {
+  test("serializeHeaderValue encodes string correctly", () => {
+    const value = "test-span-id";
+    const payload = serializeHeaderValue(value);
+
+    expect(payload.metadata?.encoding).toBeDefined();
+    expect(payload.data).toBeDefined();
+    expect(new TextDecoder().decode(payload.metadata?.encoding)).toBe(
+      "json/plain",
+    );
+    expect(new TextDecoder().decode(payload.data)).toBe('"test-span-id"');
+  });
+
+  test("deserializeHeaderValue decodes payload correctly", () => {
+    const original = "test-value-123";
+    const payload = serializeHeaderValue(original);
+    const decoded = deserializeHeaderValue(payload);
+
+    expect(decoded).toBe(original);
+  });
+
+  test("deserializeHeaderValue handles undefined payload", () => {
+    expect(deserializeHeaderValue(undefined)).toBeUndefined();
+  });
+
+  test("deserializeHeaderValue handles payload without data", () => {
+    expect(deserializeHeaderValue({ metadata: {} })).toBeUndefined();
+  });
+
+  test("deserializeHeaderValue handles invalid JSON", () => {
+    const payload = {
+      data: new TextEncoder().encode("not valid json"),
+    };
+    expect(deserializeHeaderValue(payload)).toBeUndefined();
+  });
+
+  test("round-trip serialization preserves complex strings", () => {
+    const testCases = [
+      "simple",
+      "with spaces",
+      "with/slashes",
+      "unicode-日本語",
+      "emoji-👋",
+      "",
+    ];
+
+    for (const value of testCases) {
+      const payload = serializeHeaderValue(value);
+      const decoded = deserializeHeaderValue(payload);
+      expect(decoded).toBe(value);
+    }
+  });
+
+  test("header constants are defined", () => {
+    expect(BRAINTRUST_SPAN_HEADER).toBe("_braintrust-span");
+    expect(BRAINTRUST_WORKFLOW_SPAN_HEADER).toBe("_braintrust-workflow-span");
+    expect(BRAINTRUST_WORKFLOW_SPAN_ID_HEADER).toBe(
+      "_braintrust-workflow-span-id",
+    );
+  });
+});
+
+describe("SpanComponentsV3 cross-worker reconstruction", () => {
+  test("can parse and reconstruct span components with new span_id", () => {
+    const clientComponents = new SpanComponentsV3({
+      object_type: SpanObjectTypeV3.PROJECT_LOGS,
+      object_id: "project-123",
+      row_id: "row-456",
+      span_id: "client-span-id",
+      root_span_id: "root-span-id",
+    });
+
+    const clientContext = clientComponents.toStr();
+    const workflowSpanId = "workflow-span-id";
+
+    const parsed = SpanComponentsV3.fromStr(clientContext);
+    const data = parsed.data;
+
+    expect(data.row_id).toBe("row-456");
+    expect(data.root_span_id).toBe("root-span-id");
+    expect(data.span_id).toBe("client-span-id");
+
+    if (data.row_id && data.root_span_id) {
+      const workflowComponents = new SpanComponentsV3({
+        object_type: data.object_type,
+        object_id: data.object_id,
+        propagated_event: data.propagated_event,
+        row_id: data.row_id,
+        root_span_id: data.root_span_id,
+        span_id: workflowSpanId,
+      });
+
+      const reconstructed = SpanComponentsV3.fromStr(
+        workflowComponents.toStr(),
+      );
+      expect(reconstructed.data.span_id).toBe(workflowSpanId);
+      expect(reconstructed.data.row_id).toBe("row-456");
+      expect(reconstructed.data.root_span_id).toBe("root-span-id");
+      expect(reconstructed.data.object_id).toBe("project-123");
+    }
+  });
+
+  test("preserves object_type when reconstructing", () => {
+    const objectTypes = [
+      SpanObjectTypeV3.PROJECT_LOGS,
+      SpanObjectTypeV3.EXPERIMENT,
+      SpanObjectTypeV3.PLAYGROUND_LOGS,
+    ];
+
+    for (const objectType of objectTypes) {
+      const original = new SpanComponentsV3({
+        object_type: objectType,
+        object_id: "test-id",
+        row_id: "row-id",
+        span_id: "original-span-id",
+        root_span_id: "root-span-id",
+      });
+
+      const parsed = SpanComponentsV3.fromStr(original.toStr());
+
+      const reconstructed = new SpanComponentsV3({
+        object_type: parsed.data.object_type,
+        object_id: parsed.data.object_id,
+        propagated_event: parsed.data.propagated_event,
+        row_id: parsed.data.row_id!,
+        root_span_id: parsed.data.root_span_id!,
+        span_id: "new-span-id",
+      });
+
+      expect(reconstructed.data.object_type).toBe(objectType);
+    }
+  });
+
+  test("handles span components without row_id fields", () => {
+    const componentsWithoutRowId = new SpanComponentsV3({
+      object_type: SpanObjectTypeV3.PROJECT_LOGS,
+      object_id: "test-id",
+    });
+
+    const parsed = SpanComponentsV3.fromStr(componentsWithoutRowId.toStr());
+
+    expect(parsed.data.row_id).toBeUndefined();
+    expect(parsed.data.span_id).toBeUndefined();
+    expect(parsed.data.root_span_id).toBeUndefined();
+  });
+
+  test("preserves propagated_event when reconstructing", () => {
+    const propagatedEvent = { key: "value", nested: { inner: 123 } };
+
+    const original = new SpanComponentsV3({
+      object_type: SpanObjectTypeV3.PROJECT_LOGS,
+      object_id: "test-id",
+      propagated_event: propagatedEvent,
+      row_id: "row-id",
+      span_id: "original-span-id",
+      root_span_id: "root-span-id",
+    });
+
+    const parsed = SpanComponentsV3.fromStr(original.toStr());
+
+    const reconstructed = new SpanComponentsV3({
+      object_type: parsed.data.object_type,
+      object_id: parsed.data.object_id,
+      propagated_event: parsed.data.propagated_event,
+      row_id: parsed.data.row_id!,
+      root_span_id: parsed.data.root_span_id!,
+      span_id: "new-span-id",
+    });
+
+    expect(reconstructed.data.propagated_event).toEqual(propagatedEvent);
+  });
+});
+
+describe("BraintrustTemporalPlugin", () => {
+  test("createBraintrustTemporalPlugin returns a plugin instance", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    expect(plugin).toBeInstanceOf(BraintrustTemporalPlugin);
+    expect(plugin.name).toBe("braintrust");
+  });
+
+  test("plugin has configureClient method", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    expect(typeof plugin.configureClient).toBe("function");
+  });
+
+  test("plugin has configureWorker method", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    expect(typeof plugin.configureWorker).toBe("function");
+  });
+
+  test("configureClient adds workflow interceptor", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    const options = {};
+    const configured = plugin.configureClient(options);
+
+    expect(configured.interceptors).toBeDefined();
+    expect(configured.interceptors?.workflow).toBeDefined();
+    expect(Array.isArray(configured.interceptors?.workflow)).toBe(true);
+    expect(configured.interceptors?.workflow?.length).toBe(1);
+  });
+
+  test("configureClient preserves existing interceptors", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    const existingInterceptor = { start: async (i: unknown, n: unknown) => n };
+    const options = {
+      interceptors: {
+        workflow: [existingInterceptor],
+      },
+    };
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const configured = plugin.configureClient(options as any);
+
+    expect(configured.interceptors?.workflow?.length).toBe(2);
+  });
+
+  test("configureWorker adds activity interceptor and sinks", () => {
+    const plugin = createBraintrustTemporalPlugin();
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const options = {} as any;
+    const configured = plugin.configureWorker(options);
+
+    expect(configured.interceptors).toBeDefined();
+    expect(configured.interceptors?.activity).toBeDefined();
+    expect(Array.isArray(configured.interceptors?.activity)).toBe(true);
+    expect(configured.interceptors?.activity?.length).toBe(1);
+    expect(configured.sinks).toBeDefined();
+    expect(configured.sinks?.braintrust).toBeDefined();
+  });
+});

package/src/utils.ts

@@ -0,0 +1,29 @@
+import type { Payload } from "@temporalio/common";
+
+export const BRAINTRUST_SPAN_HEADER = "_braintrust-span";
+export const BRAINTRUST_WORKFLOW_SPAN_HEADER = "_braintrust-workflow-span";
+export const BRAINTRUST_WORKFLOW_SPAN_ID_HEADER =
+  "_braintrust-workflow-span-id";
+
+export function serializeHeaderValue(value: string): Payload {
+  return {
+    metadata: {
+      encoding: new TextEncoder().encode("json/plain"),
+    },
+    data: new TextEncoder().encode(JSON.stringify(value)),
+  };
+}
+
+export function deserializeHeaderValue(
+  payload: Payload | undefined,
+): string | undefined {
+  if (!payload?.data) {
+    return undefined;
+  }
+  try {
+    const decoded = new TextDecoder().decode(payload.data);
+    return JSON.parse(decoded);
+  } catch {
+    return undefined;
+  }
+}

package/src/workflow-interceptors.ts

@@ -0,0 +1,159 @@
+/**
+ * Workflow interceptors for Braintrust tracing.
+ *
+ * IMPORTANT: This module is loaded into the Temporal workflow isolate.
+ * It cannot import Node.js modules or access external state directly.
+ * Communication with the outside world is done via sinks.
+ */
+import {
+  WorkflowInterceptorsFactory,
+  WorkflowInboundCallsInterceptor,
+  WorkflowOutboundCallsInterceptor,
+  WorkflowExecuteInput,
+  Next,
+  proxySinks,
+  workflowInfo,
+  uuid4,
+} from "@temporalio/workflow";
+import type {
+  ActivityInput,
+  LocalActivityInput,
+  StartChildWorkflowExecutionInput,
+} from "@temporalio/workflow";
+import type { Payload } from "@temporalio/common";
+import type { BraintrustSinks } from "./sinks";
+import {
+  BRAINTRUST_SPAN_HEADER,
+  BRAINTRUST_WORKFLOW_SPAN_HEADER,
+  BRAINTRUST_WORKFLOW_SPAN_ID_HEADER,
+  serializeHeaderValue,
+  deserializeHeaderValue,
+} from "./utils";
+
+const { braintrust } = proxySinks<BraintrustSinks>();
+
+/**
+ * Shared state between inbound and outbound interceptors for a single workflow.
+ * Created per-workflow by the factory function to avoid global state issues.
+ */
+interface WorkflowSpanState {
+  parentContext: string | undefined;
+  spanId: string | undefined;
+}
+
+class BraintrustWorkflowInboundInterceptor
+  implements WorkflowInboundCallsInterceptor
+{
+  constructor(private state: WorkflowSpanState) {}
+
+  async execute(
+    input: WorkflowExecuteInput,
+    next: Next<WorkflowInboundCallsInterceptor, "execute">,
+  ): Promise<unknown> {
+    // Extract parent context from headers
+    const parentContext = input.headers
+      ? deserializeHeaderValue(input.headers[BRAINTRUST_SPAN_HEADER])
+      : undefined;
+
+    // Store for the outbound interceptor to forward to activities
+    this.state.parentContext = parentContext;
+
+    // Generate a deterministic spanId for the workflow span
+    this.state.spanId = uuid4();
+
+    // Create workflow span via sink (only called if not replaying)
+    // NOTE: WorkflowInfo is injected automatically by the runtime
+    braintrust.workflowStarted(parentContext, this.state.spanId);
+
+    try {
+      const result = await next(input);
+      braintrust.workflowCompleted();
+      return result;
+    } catch (e) {
+      braintrust.workflowCompleted(e instanceof Error ? e.message : String(e));
+      throw e;
+    }
+  }
+}
+
+class BraintrustWorkflowOutboundInterceptor
+  implements WorkflowOutboundCallsInterceptor
+{
+  constructor(private state: WorkflowSpanState) {}
+
+  private getHeaders(): Record<string, Payload> {
+    const info = workflowInfo();
+    const headers: Record<string, Payload> = {};
+
+    // Pass runId so activity can look up workflow span on same worker
+    headers[BRAINTRUST_WORKFLOW_SPAN_HEADER] = serializeHeaderValue(info.runId);
+
+    // Pass workflow span ID for cross-worker activities to construct parent
+    if (this.state.spanId) {
+      headers[BRAINTRUST_WORKFLOW_SPAN_ID_HEADER] = serializeHeaderValue(
+        this.state.spanId,
+      );
+    }
+
+    // Pass client context for cross-worker activities to construct parent
+    if (this.state.parentContext) {
+      headers[BRAINTRUST_SPAN_HEADER] = serializeHeaderValue(
+        this.state.parentContext,
+      );
+    }
+
+    return headers;
+  }
+
+  scheduleActivity(
+    input: ActivityInput,
+    next: Next<WorkflowOutboundCallsInterceptor, "scheduleActivity">,
+  ) {
+    return next({
+      ...input,
+      headers: {
+        ...input.headers,
+        ...this.getHeaders(),
+      },
+    });
+  }
+
+  scheduleLocalActivity(
+    input: LocalActivityInput,
+    next: Next<WorkflowOutboundCallsInterceptor, "scheduleLocalActivity">,
+  ) {
+    return next({
+      ...input,
+      headers: {
+        ...input.headers,
+        ...this.getHeaders(),
+      },
+    });
+  }
+
+  startChildWorkflowExecution(
+    input: StartChildWorkflowExecutionInput,
+    next: Next<WorkflowOutboundCallsInterceptor, "startChildWorkflowExecution">,
+  ) {
+    return next({
+      ...input,
+      headers: {
+        ...input.headers,
+        ...this.getHeaders(),
+      },
+    });
+  }
+}
+
+export const interceptors: WorkflowInterceptorsFactory = () => {
+  // Create shared state for this workflow instance
+  const state: WorkflowSpanState = {
+    parentContext: undefined,
+    spanId: undefined,
+  };
+
+  return {
+    inbound: [new BraintrustWorkflowInboundInterceptor(state)],
+    outbound: [new BraintrustWorkflowOutboundInterceptor(state)],
+  };
+};

package/tsconfig.json

@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "declaration": true,
+    "lib": ["es2022"],
+    "module": "commonjs",
+    "target": "es2022",
+    "moduleResolution": "node",
+    "strict": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true
+  },
+  "include": ["."],
+  "exclude": ["node_modules/**", "**/dist/**"]
+}

package/tsup.config.ts

@@ -0,0 +1,20 @@
+import { defineConfig } from "tsup";
+
+export default defineConfig([
+  {
+    entry: ["src/index.ts", "src/workflow-interceptors.ts"],
+    format: ["cjs", "esm"],
+    outDir: "dist",
+    external: [
+      "braintrust",
+      "braintrust/util",
+      "@braintrust/temporal/workflow-interceptors",
+      "@temporalio/activity",
+      "@temporalio/client",
+      "@temporalio/common",
+      "@temporalio/worker",
+      "@temporalio/workflow",
+    ],
+    dts: true,
+  },
+]);

package/turbo.json

@@ -0,0 +1,8 @@
+{
+  "extends": ["//"],
+  "tasks": {
+    "build": {
+      "outputs": ["**/dist/**"]
+    }
+  }
+}

package/vitest.config.ts

@@ -0,0 +1,5 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  test: {},
+});