@braintrust/pi-extension 0.1.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Braintrust
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,126 @@
+# @braintrust/pi-extension
+
+Braintrust extension for [pi](https://github.com/mariozechner/pi-coding-agent).
+
+Today this extension automatically traces pi sessions, turns, model calls, and tool executions to Braintrust.
+
+## What gets traced
+
+- **Session spans**: one root span per pi session that actually produces at least one turn
+- **Turn spans**: one span per user prompt / agent run
+- **LLM spans**: one span per model response inside a turn
+- **Tool spans**: one span per tool execution
+
+Trace shape:
+
+```text
+Session (task)
+├── Turn 1 (task)
+│   ├── anthropic/claude-sonnet-4 (llm)
+│   │   ├── read: package.json (tool)
+│   │   └── bash: pnpm test (tool)
+│   └── anthropic/claude-sonnet-4 (llm)
+└── Turn 2 (task)
+```
+
+## Install
+
+### From npm
+
+```bash
+pi install npm:@braintrust/pi-extension
+```
+
+### From this repo
+
+```bash
+pi install .
+```
+
+Or load it just for one run:
+
+```bash
+pi -e .
+```
+
+## Compatibility
+
+This package supports the **last three stable pi versions**.
+
+Our GitHub Actions compatibility job automatically resolves and tests the latest patch release from each of the last three stable pi minor versions, so new pi releases are picked up without manually updating the matrix.
+
+## Quick start
+
+Tracing is disabled by default.
+
+Set these environment variables:
+
+```bash
+export TRACE_TO_BRAINTRUST=true
+export BRAINTRUST_API_KEY=sk-...
+export BRAINTRUST_PROJECT=pi
+```
+
+Then start pi normally.
+
+In interactive mode, the footer shows a `Braintrust` status indicator while tracing is active, and a widget below the editor shows a shortened clickable trace link when available.
+
+## Configuration
+
+You can configure the extension with environment variables or JSON config files.
+
+Config precedence is:
+
+1. defaults
+2. `~/.pi/agent/braintrust.json`
+3. `.pi/braintrust.json`
+4. environment variables
+
+### Config file locations
+
+- Global: `~/.pi/agent/braintrust.json`
+- Project: `.pi/braintrust.json`
+
+Example:
+
+```json
+{
+  "trace_to_braintrust": true,
+  "project": "pi",
+  "debug": true,
+  "additional_metadata": {
+    "team": "platform"
+  }
+}
+```
+
+## Supported settings
+
+| Config key | Env var | Default |
+|---|---|---|
+| `trace_to_braintrust` | `TRACE_TO_BRAINTRUST` | `false` |
+| `api_key` | `BRAINTRUST_API_KEY` | unset |
+| `api_url` | `BRAINTRUST_API_URL` | `https://api.braintrust.dev` |
+| `app_url` | `BRAINTRUST_APP_URL` | `https://www.braintrust.dev` |
+| `org_name` | `BRAINTRUST_ORG_NAME` | unset |
+| `project` | `BRAINTRUST_PROJECT` | `pi` |
+| `debug` | `BRAINTRUST_DEBUG` | `false` |
+| `additional_metadata` | `BRAINTRUST_ADDITIONAL_METADATA` | `{}` |
+| `log_file` | `BRAINTRUST_LOG_FILE` | unset |
+| `state_dir` | `BRAINTRUST_STATE_DIR` | `~/.pi/agent/state/braintrust-pi-extension` |
+| `parent_span_id` | `PI_PARENT_SPAN_ID` | unset |
+| `root_span_id` | `PI_ROOT_SPAN_ID` | unset |
+
+## Notes
+
+- Project config overrides global config.
+- Environment variables override both config files.
+- Session bookkeeping is stored in `~/.pi/agent/state/braintrust-pi-extension/` by default.
+- Span delivery uses the Braintrust JavaScript SDK's built-in async/background flushing.
+- If Braintrust is unavailable, pi should continue working normally.
+- If `PI_PARENT_SPAN_ID` is set, the pi session span is attached under an existing Braintrust trace.
+- `PI_ROOT_SPAN_ID` can be used when the parent span is not the trace root.
+
+## Contributing
+
+See [CONTRIBUTING.md](./CONTRIBUTING.md) for development setup, validation, and repository conventions.

package/package.json

@@ -0,0 +1,71 @@
+{
+  "name": "@braintrust/pi-extension",
+  "version": "0.1.0",
+  "description": "Braintrust extension for pi. Includes automatic tracing for pi sessions, turns, LLM calls, and tool executions to Braintrust.",
+  "keywords": [
+    "braintrust",
+    "llm",
+    "observability",
+    "pi",
+    "pi-package",
+    "tracing"
+  ],
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/braintrustdata/braintrust-pi-extension"
+  },
+  "files": [
+    "src",
+    "README.md"
+  ],
+  "type": "module",
+  "publishConfig": {
+    "access": "public"
+  },
+  "dependencies": {
+    "braintrust": "^3.7.0",
+    "valibot": "^1.3.1"
+  },
+  "devDependencies": {
+    "@mariozechner/pi-ai": "^0.64.0",
+    "@mariozechner/pi-coding-agent": "^0.64.0",
+    "@types/node": "^25.5.0",
+    "tsx": "^4.21.0",
+    "typescript": "^6.0.2",
+    "vite-plus": "^0.1.14",
+    "vitest": "^4.1.2"
+  },
+  "peerDependencies": {
+    "@mariozechner/pi-coding-agent": "*"
+  },
+  "devEngines": {
+    "packageManager": {
+      "name": "pnpm",
+      "version": ">=10.33.0",
+      "onFail": "error"
+    }
+  },
+  "engines": {
+    "npm": "please-use-pnpm",
+    "pnpm": ">=10.33.0",
+    "yarn": "please-use-pnpm"
+  },
+  "pi": {
+    "extensions": [
+      "./src/index.ts"
+    ]
+  },
+  "scripts": {
+    "preinstall": "node -e \"const userAgent = process.env.npm_config_user_agent || ''; if (process.env.INIT_CWD === process.cwd() && !userAgent.includes('pnpm/')) { console.error('Use pnpm in this repo.'); process.exit(1); }\"",
+    "check": "vp check",
+    "fmt": "vp fmt",
+    "lint": "vp lint",
+    "pack": "vp pack",
+    "test": "vitest run",
+    "test:integration": "vitest run src/index.integration.test.ts",
+    "test:watch": "vitest",
+    "typecheck": "vp check",
+    "smoke": "tsx -e \"import('./src/index.ts')\""
+  }
+}

package/src/client.ts

@@ -0,0 +1,244 @@
+import {
+  ERR_PERMALINK,
+  NOOP_SPAN_PERMALINK,
+  initLogger,
+  type Logger as BraintrustSdkLogger,
+  type Span as BraintrustSdkSpan,
+} from "braintrust";
+import type { Logger, TraceConfig } from "./types.ts";
+import { toUnixSeconds } from "./utils.ts";
+
+export type BraintrustSpanHandle = BraintrustSdkSpan;
+
+export interface StartTraceSpanArgs {
+  spanId: string;
+  rootSpanId: string;
+  parentSpanId?: string;
+  name: string;
+  type: "task" | "llm" | "tool";
+  startedAt: number;
+  input?: unknown;
+  output?: unknown;
+  error?: string;
+  metadata?: Record<string, unknown>;
+  metrics?: Record<string, number | undefined>;
+}
+
+export interface UpdateTraceSpanArgs {
+  id: string;
+  spanId?: string;
+  rootSpanId?: string;
+  input?: unknown;
+  output?: unknown;
+  error?: string;
+  metadata?: Record<string, unknown>;
+  metrics?: Record<string, number | undefined>;
+}
+
+function compactRecord<T extends Record<string, unknown>>(value: T): Partial<T> {
+  return Object.fromEntries(
+    Object.entries(value).filter(([, entry]) => entry !== undefined),
+  ) as Partial<T>;
+}
+
+function isUsablePermalink(url: string | undefined): url is string {
+  return Boolean(url && url !== NOOP_SPAN_PERMALINK && !url.startsWith(ERR_PERMALINK));
+}
+
+export class BraintrustClient {
+  readonly config: TraceConfig;
+  readonly logger?: Logger;
+  sdkLogger?: BraintrustSdkLogger<true>;
+  initPromise?: Promise<void>;
+
+  constructor(config: TraceConfig, logger?: Logger) {
+    this.config = config;
+    this.logger = logger;
+  }
+
+  #ensureLogger(): BraintrustSdkLogger<true> {
+    if (this.sdkLogger) return this.sdkLogger;
+
+    if (this.config.apiUrl && !process.env.BRAINTRUST_API_URL) {
+      process.env.BRAINTRUST_API_URL = this.config.apiUrl;
+    }
+
+    this.sdkLogger = initLogger({
+      projectName: this.config.projectName,
+      apiKey: this.config.apiKey,
+      appUrl: this.config.appUrl,
+      orgName: this.config.orgName,
+      asyncFlush: true,
+      setCurrent: false,
+      debugLogLevel: this.config.debug ? "debug" : false,
+    });
+
+    return this.sdkLogger;
+  }
+
+  async initialize(): Promise<void> {
+    if (!this.initPromise) {
+      this.initPromise = this.#doInitialize();
+    }
+    return this.initPromise;
+  }
+
+  async #doInitialize(): Promise<void> {
+    if (!this.config.apiKey) {
+      throw new Error("BRAINTRUST_API_KEY is not set");
+    }
+
+    const sdkLogger = this.#ensureLogger();
+    await sdkLogger.id;
+    const project = await sdkLogger.project;
+    this.logger?.info("braintrust sdk logger initialized", {
+      appUrl: this.config.appUrl,
+      apiUrl: this.config.apiUrl,
+      project: project.name,
+      projectId: project.id,
+    });
+  }
+
+  startSpan(args: StartTraceSpanArgs): BraintrustSpanHandle | undefined {
+    try {
+      const sdkLogger = this.#ensureLogger();
+      const span = sdkLogger.startSpan({
+        spanId: args.spanId,
+        parentSpanIds: args.parentSpanId
+          ? {
+              spanId: args.parentSpanId,
+              rootSpanId: args.rootSpanId,
+            }
+          : undefined,
+        name: args.name,
+        type: args.type,
+        startTime: toUnixSeconds(args.startedAt),
+        event: compactRecord({
+          input: args.input,
+          output: args.output,
+          error: args.error,
+          metadata: args.metadata,
+          metrics: args.metrics,
+        }),
+      });
+      return span;
+    } catch (error) {
+      this.logger?.error("failed to start Braintrust span", {
+        error: String(error),
+        spanId: args.spanId,
+        parentSpanId: args.parentSpanId,
+        rootSpanId: args.rootSpanId,
+        name: args.name,
+        type: args.type,
+      });
+      return undefined;
+    }
+  }
+
+  logSpan(span: BraintrustSpanHandle | undefined, event: Omit<UpdateTraceSpanArgs, "id">): void {
+    if (!span) return;
+
+    try {
+      span.log(
+        compactRecord({
+          input: event.input,
+          output: event.output,
+          error: event.error,
+          metadata: event.metadata,
+          metrics: event.metrics,
+        }),
+      );
+    } catch (error) {
+      this.logger?.error("failed to log Braintrust span", {
+        error: String(error),
+        spanId: span.spanId,
+        rootSpanId: span.rootSpanId,
+      });
+    }
+  }
+
+  endSpan(span: BraintrustSpanHandle | undefined, endedAt = Date.now()): void {
+    if (!span) return;
+
+    try {
+      span.end({ endTime: toUnixSeconds(endedAt) });
+    } catch (error) {
+      this.logger?.error("failed to end Braintrust span", {
+        error: String(error),
+        spanId: span.spanId,
+        rootSpanId: span.rootSpanId,
+      });
+    }
+  }
+
+  getSpanLink(span: BraintrustSpanHandle | undefined): string | undefined {
+    if (!span) return undefined;
+
+    try {
+      const link = span.link();
+      return isUsablePermalink(link) ? link : undefined;
+    } catch (error) {
+      this.logger?.warn("failed to build Braintrust span link", {
+        error: String(error),
+        spanId: span.spanId,
+        rootSpanId: span.rootSpanId,
+      });
+      return undefined;
+    }
+  }
+
+  async getSpanPermalink(span: BraintrustSpanHandle | undefined): Promise<string | undefined> {
+    if (!span) return undefined;
+
+    const link = this.getSpanLink(span);
+    if (link) return link;
+
+    try {
+      const permalink = await span.permalink();
+      return isUsablePermalink(permalink) ? permalink : undefined;
+    } catch (error) {
+      this.logger?.warn("failed to build Braintrust span permalink", {
+        error: String(error),
+        spanId: span.spanId,
+        rootSpanId: span.rootSpanId,
+      });
+      return undefined;
+    }
+  }
+
+  updateSpan(args: UpdateTraceSpanArgs): void {
+    try {
+      const sdkLogger = this.#ensureLogger();
+      sdkLogger.updateSpan({
+        id: args.id,
+        span_id: args.spanId,
+        root_span_id: args.rootSpanId,
+        input: args.input,
+        output: args.output,
+        error: args.error,
+        metadata: args.metadata,
+        metrics: args.metrics,
+      });
+    } catch (error) {
+      this.logger?.error("failed to update Braintrust span", {
+        error: String(error),
+        id: args.id,
+        spanId: args.spanId,
+        rootSpanId: args.rootSpanId,
+      });
+    }
+  }
+
+  async flush(): Promise<void> {
+    const sdkLogger = this.sdkLogger;
+    if (!sdkLogger) return;
+
+    try {
+      await sdkLogger.flush();
+    } catch (error) {
+      this.logger?.error("failed to flush Braintrust logs", {
+        error: String(error),
+      });
+    }
+  }
+}