@braintrust/pi-extension 0.1.0

package/src/state.test.ts

@@ -0,0 +1,131 @@
+import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+import { createStateStore } from "./state.ts";
+
+const tempDirs: string[] = [];
+
+afterEach(() => {
+  while (tempDirs.length > 0) {
+    rmSync(tempDirs.pop()!, { recursive: true, force: true });
+  }
+});
+
+function makeTempDir(prefix: string): string {
+  const dir = mkdtempSync(join(tmpdir(), prefix));
+  tempDirs.push(dir);
+  return dir;
+}
+
+describe("createStateStore", () => {
+  it("loads valid sessions and prunes expired ones on startup", async () => {
+    const stateDir = makeTempDir("pi-extension-state-");
+    const now = Date.now();
+    const oldTimestamp = now - 31 * 24 * 60 * 60 * 1000;
+
+    writeFileSync(
+      join(stateDir, "sessions.json"),
+      `${JSON.stringify(
+        {
+          version: 1,
+          sessions: {
+            fresh: {
+              rootSpanId: "fresh-root",
+              startedAt: now,
+              totalTurns: 2,
+            },
+            expired: {
+              rootSpanId: "old-root",
+              startedAt: oldTimestamp,
+              lastSeenAt: oldTimestamp,
+            },
+            invalid: {
+              startedAt: now,
+            },
+          },
+        },
+        null,
+        2,
+      )}\n`,
+      "utf8",
+    );
+
+    const store = createStateStore(stateDir);
+    await store.flush();
+
+    expect(store.get("fresh")).toMatchObject({ rootSpanId: "fresh-root", totalTurns: 2 });
+    expect(store.get("expired")).toBeUndefined();
+    expect(store.get("invalid")).toBeUndefined();
+
+    const persisted = JSON.parse(readFileSync(join(stateDir, "sessions.json"), "utf8")) as {
+      sessions: Record<string, unknown>;
+    };
+    expect(Object.keys(persisted.sessions)).toEqual(["fresh"]);
+  });
+
+  it("only writes to disk when persistence is scheduled or flushed", async () => {
+    const stateDir = makeTempDir("pi-extension-state-");
+    const store = createStateStore(stateDir);
+
+    store.set("session-1", {
+      rootSpanId: "root-1",
+      startedAt: 1,
+      totalTurns: 1,
+    });
+
+    expect(existsSync(join(stateDir, "sessions.json"))).toBe(false);
+
+    store.schedulePersist(0);
+    await store.flush();
+
+    const persisted = JSON.parse(readFileSync(join(stateDir, "sessions.json"), "utf8")) as {
+      sessions: Record<string, unknown>;
+    };
+    expect(persisted.sessions["session-1"]).toEqual({
+      rootSpanId: "root-1",
+      startedAt: 1,
+      totalTurns: 1,
+    });
+  });
+
+  it("persists set, patch, and delete operations", async () => {
+    const stateDir = makeTempDir("pi-extension-state-");
+    const store = createStateStore(stateDir);
+
+    store.set("session-1", {
+      rootSpanId: "root-1",
+      startedAt: 1,
+      totalTurns: 1,
+    });
+    store.patch("session-1", {
+      totalTurns: 3,
+      totalToolCalls: 5,
+    });
+    await store.flush();
+
+    expect(store.get("session-1")).toEqual({
+      rootSpanId: "root-1",
+      startedAt: 1,
+      totalTurns: 3,
+      totalToolCalls: 5,
+    });
+
+    let persisted = JSON.parse(readFileSync(join(stateDir, "sessions.json"), "utf8")) as {
+      sessions: Record<string, unknown>;
+    };
+    expect(persisted.sessions["session-1"]).toEqual({
+      rootSpanId: "root-1",
+      startedAt: 1,
+      totalTurns: 3,
+      totalToolCalls: 5,
+    });
+
+    store.delete("session-1");
+    await store.flush();
+    persisted = JSON.parse(readFileSync(join(stateDir, "sessions.json"), "utf8")) as {
+      sessions: Record<string, unknown>;
+    };
+    expect(persisted.sessions).toEqual({});
+  });
+});

package/src/state.ts

@@ -0,0 +1,197 @@
+import { existsSync, readFileSync } from "node:fs";
+import { rename, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import type { Logger, PersistedSessionState, StateStore } from "./types.ts";
+import { ensureDir, isPlainObject } from "./utils.ts";
+
+const STATE_VERSION = 1;
+const RETENTION_MS = 30 * 24 * 60 * 60 * 1000;
+const DEFAULT_PERSIST_DEBOUNCE_MS = 50;
+
+interface StoreData {
+  version: number;
+  sessions: Record<string, PersistedSessionState>;
+}
+
+interface StoreBackend {
+  stateFile: string;
+  state: StoreData;
+  dirty: boolean;
+  scheduledPersistTimer?: ReturnType<typeof setTimeout>;
+  pendingPersist: Promise<void>;
+  lastPersistedSnapshot?: string;
+  logger?: Logger;
+}
+
+const storeBackends = new Map<string, StoreBackend>();
+
+function normalizeSessions(value: unknown): Record<string, PersistedSessionState> {
+  if (!isPlainObject(value)) return {};
+
+  const sessions: Record<string, PersistedSessionState> = {};
+  for (const [key, session] of Object.entries(value)) {
+    if (
+      isPlainObject(session) &&
+      typeof session.rootSpanId === "string" &&
+      typeof session.startedAt === "number"
+    ) {
+      sessions[key] = {
+        rootSpanId: session.rootSpanId,
+        rootSpanRecordId:
+          typeof session.rootSpanRecordId === "string" ? session.rootSpanRecordId : undefined,
+        traceRootSpanId:
+          typeof session.traceRootSpanId === "string" ? session.traceRootSpanId : undefined,
+        parentSpanId: typeof session.parentSpanId === "string" ? session.parentSpanId : undefined,
+        traceUrl: typeof session.traceUrl === "string" ? session.traceUrl : undefined,
+        startedAt: session.startedAt,
+        totalTurns: typeof session.totalTurns === "number" ? session.totalTurns : undefined,
+        totalToolCalls:
+          typeof session.totalToolCalls === "number" ? session.totalToolCalls : undefined,
+        lastSeenAt: typeof session.lastSeenAt === "number" ? session.lastSeenAt : undefined,
+        sessionFile: typeof session.sessionFile === "string" ? session.sessionFile : undefined,
+      };
+    }
+  }
+
+  return sessions;
+}
+
+function serializedState(backend: StoreBackend): string {
+  return `${JSON.stringify(backend.state, null, 2)}\n`;
+}
+
+function markDirty(backend: StoreBackend): void {
+  backend.dirty = true;
+}
+
+function pruneExpired(backend: StoreBackend): boolean {
+  const cutoff = Date.now() - RETENTION_MS;
+  let changed = false;
+
+  for (const [key, value] of Object.entries(backend.state.sessions)) {
+    if ((value.lastSeenAt ?? value.startedAt ?? 0) < cutoff) {
+      delete backend.state.sessions[key];
+      changed = true;
+    }
+  }
+
+  if (changed) markDirty(backend);
+  return changed;
+}
+
+async function persistDirty(backend: StoreBackend): Promise<void> {
+  if (!backend.dirty) {
+    await backend.pendingPersist.catch(() => {});
+    return;
+  }
+
+  const snapshot = serializedState(backend);
+  backend.dirty = false;
+
+  backend.pendingPersist = backend.pendingPersist
+    .catch(() => {})
+    .then(async () => {
+      if (snapshot === backend.lastPersistedSnapshot) return;
+
+      const tempFile = `${backend.stateFile}.tmp`;
+      try {
+        await writeFile(tempFile, snapshot, "utf8");
+        await rename(tempFile, backend.stateFile);
+        backend.lastPersistedSnapshot = snapshot;
+      } catch (error) {
+        markDirty(backend);
+        backend.logger?.warn("failed to persist state store", { error: String(error) });
+      }
+    });
+
+  await backend.pendingPersist.catch(() => {});
+  if (backend.dirty) await persistDirty(backend);
+}
+
+function schedulePersist(backend: StoreBackend, delayMs = DEFAULT_PERSIST_DEBOUNCE_MS): void {
+  if (backend.scheduledPersistTimer) clearTimeout(backend.scheduledPersistTimer);
+  backend.scheduledPersistTimer = setTimeout(() => {
+    backend.scheduledPersistTimer = undefined;
+    void persistDirty(backend);
+  }, delayMs);
+  backend.scheduledPersistTimer.unref?.();
+}
+
+function createBackend(stateFile: string, logger?: Logger): StoreBackend {
+  let state: StoreData = {
+    version: STATE_VERSION,
+    sessions: {},
+  };
+
+  if (existsSync(stateFile)) {
+    try {
+      const parsed = JSON.parse(readFileSync(stateFile, "utf8")) as unknown;
+      if (isPlainObject(parsed)) {
+        state = {
+          version: STATE_VERSION,
+          sessions: normalizeSessions(parsed.sessions),
+        };
+      }
+    } catch (error) {
+      logger?.warn("failed to load state store", { error: String(error) });
+    }
+  }
+
+  const backend: StoreBackend = {
+    stateFile,
+    state,
+    dirty: false,
+    pendingPersist: Promise.resolve(),
+    logger,
+  };
+
+  backend.lastPersistedSnapshot = serializedState(backend);
+  if (pruneExpired(backend)) schedulePersist(backend, 0);
+  return backend;
+}
+
+export function createStateStore(stateDir: string, logger?: Logger): StateStore {
+  ensureDir(stateDir);
+  const stateFile = join(stateDir, "sessions.json");
+
+  let backend = storeBackends.get(stateFile);
+  if (!backend) {
+    backend = createBackend(stateFile, logger);
+    storeBackends.set(stateFile, backend);
+  } else if (!backend.logger) {
+    backend.logger = logger;
+  }
+
+  return {
+    get(sessionKey) {
+      return backend.state.sessions[sessionKey];
+    },
+    set(sessionKey, value) {
+      backend.state.sessions[sessionKey] = value;
+      markDirty(backend);
+      return backend.state.sessions[sessionKey];
+    },
+    patch(sessionKey, patch) {
+      backend.state.sessions[sessionKey] = {
+        ...backend.state.sessions[sessionKey],
+        ...patch,
+      } as PersistedSessionState;
+      markDirty(backend);
+      return backend.state.sessions[sessionKey];
+    },
+    delete(sessionKey) {
+      delete backend.state.sessions[sessionKey];
+      markDirty(backend);
+    },
+    schedulePersist(delayMs) {
+      schedulePersist(backend, delayMs);
+    },
+    async flush() {
+      if (backend.scheduledPersistTimer) {
+        clearTimeout(backend.scheduledPersistTimer);
+        backend.scheduledPersistTimer = undefined;
+      }
+      await persistDirty(backend);
+    },
+  };
+}

package/src/types.ts

@@ -0,0 +1,179 @@
+export type JsonPrimitive = string | number | boolean | null;
+export type JsonValue = JsonPrimitive | JsonObject | JsonValue[];
+
+export interface JsonObject {
+  [key: string]: JsonValue | undefined;
+}
+
+export type LogLevel = "debug" | "info" | "warn" | "error";
+
+export interface ConfigIssue {
+  path: string;
+  message: string;
+  severity: "error" | "warning";
+}
+
+export interface TraceConfig {
+  enabled: boolean;
+  apiKey: string;
+  apiUrl?: string;
+  appUrl: string;
+  orgName?: string;
+  projectName: string;
+  debug: boolean;
+  logFile?: string;
+  stateDir: string;
+  additionalMetadata: JsonObject;
+  parentSpanId?: string;
+  rootSpanId?: string;
+  configIssues: ConfigIssue[];
+}
+
+export interface Logger {
+  filePath: string;
+  debug(message: string, data?: unknown): void;
+  info(message: string, data?: unknown): void;
+  warn(message: string, data?: unknown): void;
+  error(message: string, data?: unknown): void;
+  flush(): Promise<void>;
+}
+
+export interface PersistedSessionState {
+  rootSpanId: string;
+  rootSpanRecordId?: string;
+  traceRootSpanId?: string;
+  parentSpanId?: string;
+  traceUrl?: string;
+  startedAt: number;
+  totalTurns?: number;
+  totalToolCalls?: number;
+  lastSeenAt?: number;
+  sessionFile?: string;
+}
+
+export interface StateStore {
+  get(sessionKey: string): PersistedSessionState | undefined;
+  set(sessionKey: string, value: PersistedSessionState): PersistedSessionState;
+  patch(sessionKey: string, patch: Partial<PersistedSessionState>): PersistedSessionState;
+  delete(sessionKey: string): void;
+  schedulePersist(delayMs?: number): void;
+  flush(): Promise<void>;
+}
+
+export interface TextContentLike {
+  type: "text";
+  text?: string;
+}
+
+export interface ThinkingContentLike {
+  type: "thinking";
+  thinking?: string;
+  redacted?: boolean;
+}
+
+export interface ToolCallContentLike {
+  type: "toolCall";
+  id: string;
+  name?: string;
+  arguments?: Record<string, unknown>;
+}
+
+export interface ImageContentLike {
+  type: "image";
+  mimeType?: string;
+}
+
+export type ContentPartLike =
+  | TextContentLike
+  | ThinkingContentLike
+  | ToolCallContentLike
+  | ImageContentLike
+  | {
+      type?: string;
+      [key: string]: unknown;
+    };
+
+export interface ImageLike {
+  mimeType?: string;
+  source?: {
+    mediaType?: string;
+  };
+}
+
+export interface UsageLike {
+  input?: number;
+  output?: number;
+  cacheRead?: number;
+  cacheWrite?: number;
+  totalTokens?: number;
+}
+
+export interface UserMessageLike {
+  role: "user";
+  content?: unknown;
+}
+
+export interface AssistantMessageLike {
+  role: "assistant";
+  content?: ContentPartLike[];
+  api?: string;
+  provider?: string;
+  model?: string;
+  usage?: UsageLike;
+  stopReason?: string;
+  errorMessage?: string;
+  timestamp?: number;
+}
+
+export interface ToolResultMessageLike {
+  role: "toolResult";
+  toolCallId?: string;
+  toolName?: string;
+  content?: unknown;
+  details?: unknown;
+  isError?: boolean;
+}
+
+export type AgentMessageLike =
+  | UserMessageLike
+  | AssistantMessageLike
+  | ToolResultMessageLike
+  | {
+      role?: string;
+      content?: unknown;
+    };
+
+export interface NormalizedUserMessage {
+  role: "user";
+  content: string;
+}
+
+export interface NormalizedAssistantMessage {
+  role: "assistant";
+  content: string;
+  tool_calls?: Array<{
+    id?: string;
+    type: "function";
+    function: {
+      name?: string;
+      arguments: string;
+    };
+  }>;
+  reasoning?: Array<{
+    id: string;
+    content: string;
+  }>;
+}
+
+export interface NormalizedToolMessage {
+  role: "tool";
+  content: string;
+  tool_call_id?: string;
+  name?: string;
+  is_error: boolean;
+}
+
+export type NormalizedAgentMessage =
+  | NormalizedUserMessage
+  | NormalizedAssistantMessage
+  | NormalizedToolMessage;

package/src/utils.test.ts

@@ -0,0 +1,163 @@
+import { execFileSync } from "node:child_process";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  buildTurnInput,
+  extractErrorText,
+  formatToolSpanName,
+  normalizeAssistantMessage,
+  normalizeContextMessages,
+  normalizeToolResult,
+  repoSlugForCwd,
+  rootSpanName,
+} from "./utils.ts";
+
+const tempDirs: string[] = [];
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  while (tempDirs.length > 0) {
+    rmSync(tempDirs.pop()!, { recursive: true, force: true });
+  }
+});
+
+function makeTempDir(prefix: string): string {
+  const dir = mkdtempSync(join(tmpdir(), prefix));
+  tempDirs.push(dir);
+  return dir;
+}
+
+describe("utils", () => {
+  it("normalizes assistant messages with text, reasoning, and tool calls", () => {
+    const normalized = normalizeAssistantMessage({
+      role: "assistant",
+      content: [
+        { type: "text", text: "First line" },
+        { type: "thinking", thinking: "Plan the next step" },
+        {
+          type: "toolCall",
+          id: "call-1",
+          name: "read",
+          arguments: { path: "package.json" },
+        },
+      ],
+    });
+
+    expect(normalized).toEqual({
+      role: "assistant",
+      content: "First line",
+      reasoning: [{ id: "thinking", content: "Plan the next step" }],
+      tool_calls: [
+        {
+          id: "call-1",
+          type: "function",
+          function: {
+            name: "read",
+            arguments: '{"path":"package.json"}',
+          },
+        },
+      ],
+    });
+  });
+
+  it("normalizes mixed context messages into Braintrust-friendly shapes", () => {
+    const messages = normalizeContextMessages([
+      { role: "user", content: [{ type: "text", text: "Use the docs" }] },
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "I will inspect the config" },
+          { type: "thinking", redacted: true },
+        ],
+      },
+      {
+        role: "toolResult",
+        toolCallId: "tool-1",
+        toolName: "read",
+        isError: false,
+        content: [{ type: "text", text: "README contents" }],
+      },
+      { role: "system", content: "ignored" },
+    ]);
+
+    expect(messages).toEqual([
+      { role: "user", content: "Use the docs" },
+      {
+        role: "assistant",
+        content: "I will inspect the config",
+        reasoning: [{ id: "thinking", content: "[thinking redacted]" }],
+      },
+      {
+        role: "tool",
+        content: "README contents",
+        tool_call_id: "tool-1",
+        name: "read",
+        is_error: false,
+      },
+    ]);
+  });
+
+  it("normalizes structured tool results and extracts readable error text", () => {
+    const result = normalizeToolResult({
+      content: [
+        { type: "text", text: "command failed" },
+        { type: "thinking", redacted: true },
+      ],
+      details: { exitCode: 1, stderr: "boom" },
+      isError: true,
+    });
+
+    expect(result).toEqual({
+      content: "command failed\n[thinking redacted]",
+      details: { exitCode: 1, stderr: "boom" },
+      isError: true,
+    });
+
+    expect(
+      extractErrorText(
+        {
+          content: [{ type: "text", text: "tool exploded" }],
+        },
+        "fallback message",
+      ),
+    ).toBe("tool exploded");
+  });
+
+  it("formats tool span names and builds turn inputs", () => {
+    expect(formatToolSpanName("read", { path: "/tmp/project/package.json" })).toBe(
+      "read: package.json",
+    );
+    expect(formatToolSpanName("bash", { command: "npm    test   -- --runInBand" })).toBe(
+      "bash: npm test -- --runInBand",
+    );
+
+    expect(
+      buildTurnInput("Summarize these screenshots", [
+        { source: { mediaType: "image/png" } },
+        { mimeType: "image/jpeg" },
+      ]),
+    ).toBe("Summarize these screenshots\n[image/png]\n[image/jpeg]");
+  });
+
+  it("prefers owner/repo from git origin for the root span name", () => {
+    const repoDir = makeTempDir("pi-extension-git-");
+    execFileSync("git", ["init"], { cwd: repoDir, stdio: "ignore" });
+    execFileSync(
+      "git",
+      ["remote", "add", "origin", "git@github.com:braintrustdata/braintrust-pi-extension.git"],
+      { cwd: repoDir, stdio: "ignore" },
+    );
+
+    expect(repoSlugForCwd(repoDir)).toBe("braintrustdata/braintrust-pi-extension");
+    expect(rootSpanName(repoDir)).toBe("pi: braintrustdata/braintrust-pi-extension");
+  });
+
+  it("falls back to the cwd basename when no git origin is available", () => {
+    const dir = makeTempDir("pi-extension-no-git-");
+
+    expect(repoSlugForCwd(dir)).toBeUndefined();
+    expect(rootSpanName(dir)).toBe(`pi: ${dir.split("/").at(-1)}`);
+  });
+});