diff --git a/AGENTS.md b/AGENTS.md index 207fffb..4c7cc30 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -27,6 +27,7 @@ Create a `braintrust.json` file in one of these locations: ```json { "trace_to_braintrust": true, + "enable_tools": true, "project": "my-project", "debug": true } @@ -37,6 +38,7 @@ Create a `braintrust.json` file in one of these locations: | Config Key | Env Var | Type | Default | Description | |------------|---------|------|---------|-------------| | `trace_to_braintrust` | `TRACE_TO_BRAINTRUST` | boolean | `false` | Enable/disable tracing | +| `enable_tools` | `BRAINTRUST_OPENCODE_ENABLE_TOOLS` | boolean | `true` | Register Braintrust tools in OpenCode | | `project` | `BRAINTRUST_PROJECT` | string | `"opencode"` | Project name for traces | | `debug` | `BRAINTRUST_DEBUG` | boolean | `false` | Enable debug logging | | `api_key` | `BRAINTRUST_API_KEY` | string | | API key for authentication | diff --git a/README.md b/README.md index 811f2ab..5a6ba98 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ Create a `braintrust.json` file in one of these locations: ```json { "trace_to_braintrust": true, + "enable_tools": true, "project": "my-project", "api_key": "your-api-key", "debug": true @@ -55,6 +56,7 @@ Create a `braintrust.json` file in one of these locations: | Config Key | Env Var | Type | Default | Description | |------------|---------|------|---------|-------------| | `trace_to_braintrust` | `TRACE_TO_BRAINTRUST` | boolean | `false` | Enable/disable tracing | +| `enable_tools` | `BRAINTRUST_OPENCODE_ENABLE_TOOLS` | boolean | `true` | Register Braintrust tools in OpenCode | | `project` | `BRAINTRUST_PROJECT` | string | `"opencode"` | Project name for traces | | `debug` | `BRAINTRUST_DEBUG` | boolean | `false` | Enable debug logging | | `api_key` | `BRAINTRUST_API_KEY` | string | | API key for authentication | @@ -72,6 +74,24 @@ Configuration is loaded with the following precedence (later overrides earlier): 3. `.opencode/braintrust.json` (project config) 4. Environment variables (highest priority) +## Disabling Braintrust Tools + +Set `enable_tools` to `false` to trace OpenCode sessions without registering Braintrust-branded tools (`braintrust_query_logs`, `braintrust_list_projects`, `braintrust_log_data`, `braintrust_get_experiments`): + +```json +{ + "trace_to_braintrust": true, + "enable_tools": false, + "project": "my-project" +} +``` + +Or use the environment variable: + +```bash +BRAINTRUST_OPENCODE_ENABLE_TOOLS=false TRACE_TO_BRAINTRUST=true opencode +``` + ## Adding Dynamic Metadata Use `BRAINTRUST_ADDITIONAL_METADATA` to attach custom key-value pairs to the root span. This is useful for tagging traces in CI or linking them back to a specific run. diff --git a/src/client.test.ts b/src/client.test.ts index e65decf..1c3c819 100644 --- a/src/client.test.ts +++ b/src/client.test.ts @@ -59,6 +59,7 @@ describe("loadConfig", () => { "BRAINTRUST_ORG_NAME", "BRAINTRUST_PROJECT", "BRAINTRUST_ADDITIONAL_METADATA", + "BRAINTRUST_OPENCODE_ENABLE_TOOLS", ] beforeEach(() => { @@ -89,6 +90,7 @@ describe("loadConfig", () => { expect(config.orgName).toBeUndefined() expect(config.projectName).toBe("opencode") expect(config.tracingEnabled).toBe(false) + expect(config.enableTools).toBe(true) expect(config.debug).toBe(false) }) }) @@ -129,6 +131,18 @@ describe("loadConfig", () => { const config = loadConfig() expect(config.debug).toBe(true) }) + + it("BRAINTRUST_OPENCODE_ENABLE_TOOLS=false disables tools", () => { + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "false" + const config = loadConfig() + expect(config.enableTools).toBe(false) + }) + + it("BRAINTRUST_OPENCODE_ENABLE_TOOLS=1 enables tools", () => { + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "1" + const config = loadConfig() + expect(config.enableTools).toBe(true) + }) }) describe("pluginConfig only (from opencode.json)", () => { @@ -140,6 +154,7 @@ describe("loadConfig", () => { org_name: "test-org", project: "test-project", trace_to_braintrust: true, + enable_tools: false, debug: true, } const config = loadConfig(pluginConfig) @@ -149,6 +164,7 @@ describe("loadConfig", () => { expect(config.orgName).toBe("test-org") expect(config.projectName).toBe("test-project") expect(config.tracingEnabled).toBe(true) + expect(config.enableTools).toBe(false) expect(config.debug).toBe(true) }) @@ -167,10 +183,12 @@ describe("loadConfig", () => { it("handles pluginConfig with false booleans", () => { const pluginConfig: PluginConfig = { trace_to_braintrust: false, + enable_tools: false, debug: false, } const config = loadConfig(pluginConfig) expect(config.tracingEnabled).toBe(false) + expect(config.enableTools).toBe(false) expect(config.debug).toBe(false) }) }) @@ -204,6 +222,20 @@ describe("loadConfig", () => { expect(config.debug).toBe(true) }) + it("env var overrides pluginConfig for enable_tools", () => { + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "true" + const pluginConfig: PluginConfig = { enable_tools: false } + const config = loadConfig(pluginConfig) + expect(config.enableTools).toBe(true) + }) + + it("env var can disable tools when pluginConfig enables them", () => { + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "0" + const pluginConfig: PluginConfig = { enable_tools: true } + const config = loadConfig(pluginConfig) + expect(config.enableTools).toBe(false) + }) + it("pluginConfig is used when env var is not set", () => { process.env.BRAINTRUST_API_KEY = "env-key" // BRAINTRUST_PROJECT not set diff --git a/src/client.ts b/src/client.ts index f9baed1..afc8008 100644 --- a/src/client.ts +++ b/src/client.ts @@ -9,6 +9,7 @@ export interface BraintrustConfig { orgName?: string projectName: string tracingEnabled: boolean + enableTools: boolean debug: boolean logToFile?: string // Path to write NDJSON log of all plugin I/O, or "auto" for default path additionalMetadata?: Record @@ -26,6 +27,7 @@ export interface PluginConfig { org_name?: string project?: string trace_to_braintrust?: boolean + enable_tools?: boolean debug?: boolean log_to_file?: string // Path to write NDJSON log of all plugin I/O, or "true"/"auto" for default path additional_metadata?: Record @@ -108,6 +110,7 @@ export function loadConfig(pluginConfig?: PluginConfig): BraintrustConfig { orgName: undefined, projectName: "opencode", tracingEnabled: false, + enableTools: true, debug: false, logToFile: undefined, } @@ -122,6 +125,9 @@ export function loadConfig(pluginConfig?: PluginConfig): BraintrustConfig { if (pluginConfig.trace_to_braintrust !== undefined) { defaults.tracingEnabled = pluginConfig.trace_to_braintrust } + if (pluginConfig.enable_tools !== undefined) { + defaults.enableTools = pluginConfig.enable_tools + } if (pluginConfig.debug !== undefined) { defaults.debug = pluginConfig.debug } @@ -170,6 +176,9 @@ export function loadConfig(pluginConfig?: PluginConfig): BraintrustConfig { tracingEnabled: process.env.TRACE_TO_BRAINTRUST ? parseBooleanEnv(process.env.TRACE_TO_BRAINTRUST) : defaults.tracingEnabled, + enableTools: process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS + ? parseBooleanEnv(process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS) + : defaults.enableTools, debug: process.env.BRAINTRUST_DEBUG ? parseBooleanEnv(process.env.BRAINTRUST_DEBUG) : defaults.debug, diff --git a/src/index.test.ts b/src/index.test.ts new file mode 100644 index 0000000..fc61c41 --- /dev/null +++ b/src/index.test.ts @@ -0,0 +1,116 @@ +/** + * Tests for OpenCode plugin registration behavior + */ + +import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test" +import { mkdtempSync, rmSync } from "node:fs" +import { tmpdir } from "node:os" +import { join } from "node:path" +import type { PluginInput } from "@opencode-ai/plugin" + +mock.module("@opencode-ai/plugin", () => ({ + tool: Object.assign((definition: unknown) => definition, { + schema: { + string: () => ({ + optional() { + return this + }, + describe() { + return this + }, + }), + number: () => ({ + optional() { + return this + }, + describe() { + return this + }, + }), + }, + }), +})) + +function createInput(directory: string): PluginInput { + return { + directory, + worktree: directory, + project: "test-project", + client: { + app: { + log: async () => {}, + }, + }, + } as unknown as PluginInput +} + +describe("BraintrustPlugin", () => { + const originalEnv: Record = {} + const envVars = [ + "TRACE_TO_BRAINTRUST", + "BRAINTRUST_API_KEY", + "BRAINTRUST_API_URL", + "BRAINTRUST_OPENCODE_ENABLE_TOOLS", + "HOME", + ] + const originalFetch = globalThis.fetch + let directory: string + let fetchCalls: number + + beforeEach(() => { + directory = mkdtempSync(join(tmpdir(), "braintrust-opencode-plugin-")) + for (const key of envVars) { + originalEnv[key] = process.env[key] + delete process.env[key] + } + process.env.HOME = directory + process.env.TRACE_TO_BRAINTRUST = "false" + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "true" + process.env.BRAINTRUST_API_KEY = "test-api-key" + process.env.BRAINTRUST_API_URL = "https://api.example.com" + fetchCalls = 0 + globalThis.fetch = (async () => { + fetchCalls++ + return new Response(JSON.stringify({ id: "project-id", name: "opencode" }), { + status: 200, + }) + }) as unknown as typeof fetch + }) + + afterEach(() => { + rmSync(directory, { recursive: true, force: true }) + for (const key of envVars) { + if (originalEnv[key] !== undefined) { + process.env[key] = originalEnv[key] + } else { + delete process.env[key] + } + } + globalThis.fetch = originalFetch + }) + + it("registers Braintrust tools when enabled", async () => { + const { BraintrustPlugin } = await import("./index") + const hooks = await BraintrustPlugin(createInput(directory)) + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(hooks.tool).toBeDefined() + expect(Object.keys(hooks.tool ?? {}).sort()).toEqual([ + "braintrust_get_experiments", + "braintrust_list_projects", + "braintrust_log_data", + "braintrust_query_logs", + ]) + }) + + it("does not register Braintrust tools when disabled", async () => { + process.env.BRAINTRUST_OPENCODE_ENABLE_TOOLS = "false" + + const { BraintrustPlugin } = await import("./index") + const hooks = await BraintrustPlugin(createInput(directory)) + await Promise.resolve() + + expect(hooks.tool).toBeUndefined() + expect(fetchCalls).toBe(0) + }) +}) diff --git a/src/index.ts b/src/index.ts index 6bd677c..8381506 100644 --- a/src/index.ts +++ b/src/index.ts @@ -63,11 +63,12 @@ export const BraintrustPlugin: Plugin = async (input: PluginInput) => { .catch(() => {}) } - // Create Braintrust client but don't initialize yet (lazy initialization) + // Create Braintrust client only when at least one Braintrust-backed feature is enabled let btClient: BraintrustClient | undefined let _initPromise: Promise | undefined + const featuresEnabled = config.tracingEnabled || config.enableTools - if (config.apiKey) { + if (config.apiKey && featuresEnabled) { btClient = new BraintrustClient(config) // Start initialization in background, don't await _initPromise = btClient.initialize().catch((error) => { @@ -118,8 +119,8 @@ export const BraintrustPlugin: Plugin = async (input: PluginInput) => { .catch(() => {}) } - // Add Braintrust tools if client is available - if (btClient) { + // Add Braintrust tools if enabled and client is available + if (btClient && config.enableTools) { hooks.tool = createBraintrustTools(btClient) } @@ -130,11 +131,11 @@ export const BraintrustPlugin: Plugin = async (input: PluginInput) => { body: { service: "braintrust", level: "info", - message: `Logging Braintrust spans to project "${config.projectName}"`, + message: `Braintrust plugin enabled for project "${config.projectName}"`, }, }) .catch(() => {}) - } else { + } else if (featuresEnabled) { client.app .log({ body: { @@ -145,6 +146,16 @@ export const BraintrustPlugin: Plugin = async (input: PluginInput) => { }, }) .catch(() => {}) + } else { + client.app + .log({ + body: { + service: "braintrust", + level: "info", + message: "Braintrust tracing and tools are disabled.", + }, + }) + .catch(() => {}) } return hooks diff --git a/src/tracing.test.ts b/src/tracing.test.ts index 2d38004..b06a396 100644 --- a/src/tracing.test.ts +++ b/src/tracing.test.ts @@ -402,6 +402,7 @@ describe("Metric timestamps use Unix seconds", () => { appUrl: "https://www.braintrust.dev", projectName: "test-project", tracingEnabled: true, + enableTools: true, debug: false, }, ) @@ -773,6 +774,7 @@ describe("System prompt capture: production hooks", () => { appUrl: "https://www.braintrust.dev", projectName: "test-project", tracingEnabled: true, + enableTools: true, debug: false, }, )