diff --git a/README.md b/README.md index 983b3a6..4ec3c58 100644 --- a/README.md +++ b/README.md @@ -235,6 +235,7 @@ Then point your MCP client at `http://localhost:8080/mcp` using the same header/ | `BROWSERLESS_TIMEOUT` | No | `30000` | Request timeout in milliseconds | | `BROWSERLESS_MAX_RETRIES` | No | `3` | Max retry attempts for failed requests | | `BROWSERLESS_CACHE_TTL` | No | `60000` | Cache TTL in milliseconds (0 to disable) | +| `AMPLITUDE_API_KEY` | No | — | Amplitude project API key for MCP usage analytics | | `MCP_COMPLIANCE_MODE` | No | unset (full surface) | Serve the reduced, directory-compliant surface. Fails closed: any set value except `false`/`0`/`no`/`off` enables it | ## MCP Resources diff --git a/llms-install.md b/llms-install.md index feaad5d..b2493ef 100644 --- a/llms-install.md +++ b/llms-install.md @@ -54,6 +54,8 @@ If the user needs a local-only / air-gapped install, use the npm package over st } ``` +To opt in to Amplitude analytics, add `AMPLITUDE_API_KEY` with a real Amplitude project key. + 4. **Reload the MCP client.** ## Verify install diff --git a/package-lock.json b/package-lock.json index a9469b4..257a254 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,7 +10,10 @@ "hasInstallScript": true, "license": "SSPL-1.0", "dependencies": { + "@amplitude/analytics-node": "^1.5.68", + "@amplitude/mcp-analytics": "^0.3.0", "@aws-sdk/client-sqs": "^3.1053.0", + "@modelcontextprotocol/sdk": "^1.24.3", "fastmcp": "4.4.0", "ioredis": "^5.10.1", "ws": "^8.21.0", @@ -45,6 +48,52 @@ "npm": ">=11.10.0" } }, + "node_modules/@amplitude/analytics-connector": { + "version": "1.6.5", + "resolved": "https://registry.npmjs.org/@amplitude/analytics-connector/-/analytics-connector-1.6.5.tgz", + "integrity": "sha512-EcDT+E4vliT7WAeOCbyE4rNkTUdLMR9dGCETX+36U7gAROr8jARG6ObhW7TTWdoYDlUe25IllZnRKdXHRKlWJQ==", + "license": "MIT" + }, + "node_modules/@amplitude/analytics-core": { + "version": "2.54.0", + "resolved": "https://registry.npmjs.org/@amplitude/analytics-core/-/analytics-core-2.54.0.tgz", + "integrity": "sha512-xZEpG5IQvRA6qBTnG+chANEeATL9oqOvQUrqKbwl+q0gX9a6skplUMBwhEhBv2iwuDAJhDF2+G2mYf7XGuEX4Q==", + "license": "MIT", + "dependencies": { + "@amplitude/analytics-connector": "^1.6.4", + "@types/zen-observable": "0.8.3", + "safe-json-stringify": "1.2.0", + "tslib": "^2.4.1", + "zen-observable": "0.10.0" + } + }, + "node_modules/@amplitude/analytics-node": { + "version": "1.5.68", + "resolved": "https://registry.npmjs.org/@amplitude/analytics-node/-/analytics-node-1.5.68.tgz", + "integrity": "sha512-ZnE1/wvGo8msi2K8L1S7aUMOhdEJrEL6olXvq2Ds5oeZqP7m0NIwGkgplapu0Vr7ISmsOZGL+z5cUtiBrX4k6A==", + "license": "MIT", + "dependencies": { + "@amplitude/analytics-core": "2.54.0", + "tslib": "^2.4.1" + } + }, + "node_modules/@amplitude/mcp-analytics": { + "version": "0.3.0", + "resolved": "https://registry.npmjs.org/@amplitude/mcp-analytics/-/mcp-analytics-0.3.0.tgz", + "integrity": "sha512-OjWJ82iQm/5JCjcJKGGYNuCgvwYItr9LITnC9P9MF3V5t/+GRJHALBs2IrBHAGNHbW1mb2SBRs5jOBRkAcJ+eA==", + "peerDependencies": { + "@amplitude/analytics-node": ">=1.3.0", + "@modelcontextprotocol/sdk": ">=1.14.0" + }, + "peerDependenciesMeta": { + "@amplitude/analytics-node": { + "optional": false + }, + "@modelcontextprotocol/sdk": { + "optional": false + } + } + }, "node_modules/@aws-sdk/client-sqs": { "version": "3.1090.0", "resolved": "https://registry.npmjs.org/@aws-sdk/client-sqs/-/client-sqs-3.1090.0.tgz", @@ -607,12 +656,12 @@ } }, "node_modules/@modelcontextprotocol/sdk": { - "version": "1.26.0", - "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.26.0.tgz", - "integrity": "sha512-Y5RmPncpiDtTXDbLKswIJzTqu2hyBKxTNsgKqKclDbhIgg1wgtf1fRuvxgTnRfcnxtvvgbIEcqUOzZrJ6iSReg==", + "version": "1.30.0", + "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.30.0.tgz", + "integrity": "sha512-xKd8OIzlqNzcqcNumGAa6g+PW2kjD5vrpcKOnfldAUPP3j7lnqMPwlTXQm8gF+UwH72z0lqaRbjr9hqGz0eITA==", "license": "MIT", "dependencies": { - "@hono/node-server": "^1.19.9", + "@hono/node-server": "^1.19.9 || ^2.0.5", "ajv": "^8.17.1", "ajv-formats": "^3.0.1", "content-type": "^1.0.5", @@ -946,6 +995,12 @@ "@types/node": "*" } }, + "node_modules/@types/zen-observable": { + "version": "0.8.3", + "resolved": "https://registry.npmjs.org/@types/zen-observable/-/zen-observable-0.8.3.tgz", + "integrity": "sha512-fbF6oTd4sGGy0xjHPKAt+eS2CrxJ3+6gQ3FGcBoIJR2TLAyCkCyI8JqZNy+FeON0AhVgNJoUumVoZQjBFUqHkw==", + "license": "MIT" + }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.64.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.64.0.tgz", @@ -1413,16 +1468,16 @@ "license": "MIT" }, "node_modules/brace-expansion": { - "version": "5.0.7", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", - "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", + "version": "5.0.8", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.8.tgz", + "integrity": "sha512-JZyDyq3D4AUifKTPOB7DELf6XsB3WdPuNxCtob1vFXPsSXhdAiHBWJ/tJ8HAc9aH84BK+5JFZLNkJKx3G9kzQg==", "dev": true, "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" }, "engines": { - "node": "18 || 20 || >=22" + "node": "20 || >=22" } }, "node_modules/browser-stdout": { @@ -4466,6 +4521,12 @@ "node": ">= 18" } }, + "node_modules/safe-json-stringify": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/safe-json-stringify/-/safe-json-stringify-1.2.0.tgz", + "integrity": "sha512-gH8eh2nZudPQO6TytOvbxnuhYBOvDBBLW52tz5q6X58lJcd/tkmqFR+5Z9adS8aJtURSXWThWy/xJtJwixErvg==", + "license": "MIT" + }, "node_modules/safer-buffer": { "version": "2.1.2", "resolved": "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz", @@ -5440,6 +5501,12 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/zen-observable": { + "version": "0.10.0", + "resolved": "https://registry.npmjs.org/zen-observable/-/zen-observable-0.10.0.tgz", + "integrity": "sha512-iI3lT0iojZhKwT5DaFy2Ce42n3yFcLdFyOh01G7H0flMY60P8MJuVFEoJoNwXlmAyQ45GrjL6AcZmmlv8A5rbw==", + "license": "MIT" + }, "node_modules/zod": { "version": "4.4.3", "resolved": "https://registry.npmjs.org/zod/-/zod-4.4.3.tgz", diff --git a/package.json b/package.json index 48cfdfd..4556315 100644 --- a/package.json +++ b/package.json @@ -75,7 +75,10 @@ "printWidth": 80 }, "dependencies": { + "@amplitude/analytics-node": "^1.5.68", + "@amplitude/mcp-analytics": "^0.3.0", "@aws-sdk/client-sqs": "^3.1053.0", + "@modelcontextprotocol/sdk": "^1.24.3", "fastmcp": "4.4.0", "ioredis": "^5.10.1", "ws": "^8.21.0", diff --git a/src/@types/types.d.ts b/src/@types/types.d.ts index 1f6e9ee..685ce23 100644 --- a/src/@types/types.d.ts +++ b/src/@types/types.d.ts @@ -44,6 +44,8 @@ export interface BrowserlessSession extends Record { * resulting id instead of letting the model open a `createProfile` session. */ attachSessionId?: string; + /** Verified account id for OAuth/Supabase-authenticated sessions. */ + accountId?: string; /** Origin tag from the `x-browserless-mcp-source` header; see resolveMcpSource. */ source?: string; } @@ -69,6 +71,7 @@ export interface McpConfig { maxRetries: number; cacheTtlMs: number; analyticsEnabled: boolean; + amplitudeApiKey?: string; // Required (not optional): the compliant surface is a security gate, so every // McpConfig must choose explicitly — an omitted field must not default to the // fuller/prohibited surface. diff --git a/src/config.ts b/src/config.ts index 051a16d..ce0ac7d 100644 --- a/src/config.ts +++ b/src/config.ts @@ -52,6 +52,7 @@ export function getConfig(): McpConfig { maxRetries: parseInt(process.env.BROWSERLESS_MAX_RETRIES ?? '3', 10), cacheTtlMs: parseInt(process.env.BROWSERLESS_CACHE_TTL ?? '60000', 10), analyticsEnabled: process.env.ANALYTICS_ENABLED === 'true', + amplitudeApiKey: process.env.AMPLITUDE_API_KEY, // Per-process toggle for the compliant surface used by the OpenAI/Anthropic // directory listings: registers fewer tools and de-fangs the agent (see // tools/compliance.ts). Fails closed — see parseComplianceMode. diff --git a/src/index.ts b/src/index.ts index 2daff21..e7c62ff 100644 --- a/src/index.ts +++ b/src/index.ts @@ -16,6 +16,12 @@ import { resolveBrowserlessAuth } from './lib/http-auth.js'; import { BoundedEventStore } from './lib/bounded-event-store.js'; import { RedisOAuthProxy } from './lib/redis-oauth-proxy.js'; import { Redis } from 'ioredis'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { + instrumentFastMcpTools, + initializeAmplitudeAnalytics, + shutdownAmplitudeAnalytics, +} from './lib/amplitude-analytics.js'; const pkg = JSON.parse( readFileSync( @@ -38,6 +44,10 @@ const analytics = new AnalyticsHelper( config.sqsQueueUrl, config.sqsRegion, ); +const amplitudeAnalytics = initializeAmplitudeAnalytics( + config.amplitudeApiKey, + pkg.version, +); // Passthrough OAuth provider: disables FastMCP's token-swap mode so the MCP client // receives the raw Supabase JWT directly. @@ -122,6 +132,7 @@ const server = new FastMCP({ authenticate: hybridAuthenticate, }); +instrumentFastMcpTools(server, amplitudeAnalytics); registerSurface(server, config, analytics); // Log the active surface (both transports) so it's visible in the boot logs. // Fail-closed value lands on compliant; distinguish "unset" (dropped/wrong-scoped @@ -143,15 +154,43 @@ const complianceSurface = config.complianceMode : 'full (explicit opt-out)'; console.error(`[browserless-mcp] Tool surface: ${complianceSurface}`); +let warnedAboutServerIdentity = false; server.on('connect', (event) => { const id = event.session.sessionId ?? 'stdio'; console.error(`[browserless-mcp] Client connected: ${id}`); + if ( + amplitudeAnalytics && + !warnedAboutServerIdentity && + !(event.session.server instanceof Server) + ) { + warnedAboutServerIdentity = true; + console.error( + '[browserless-mcp] WARNING: FastMCP session server is not an MCP SDK Server; Amplitude instrumentation may be disabled.', + ); + } // force the client to refresh its tool list on connect void event.session.triggerListChangedNotification( 'notifications/tools/list_changed', ); }); +if (amplitudeAnalytics) { + let amplitudeShutdown = false; + const shutdown = (exitCode: number): void => { + if (amplitudeShutdown) return; + amplitudeShutdown = true; + void (async () => { + try { + await shutdownAmplitudeAnalytics(amplitudeAnalytics); + } finally { + process.exit(exitCode); + } + })(); + }; + process.once('SIGTERM', () => shutdown(143)); + process.once('SIGINT', () => shutdown(130)); +} + server.on('disconnect', (event) => { const id = event.session.sessionId ?? 'stdio'; // Drop any files staged/captured for this session (TTL is the backstop). diff --git a/src/lib/account-resolver.ts b/src/lib/account-resolver.ts index e551ee1..c223a07 100644 --- a/src/lib/account-resolver.ts +++ b/src/lib/account-resolver.ts @@ -5,6 +5,7 @@ import type { SupabaseJwtPayload } from '../@types/types.js'; interface ResolvedAccount { apiKey: string; email: string; + accountId: string; } const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes @@ -149,6 +150,7 @@ export async function resolveApiKey( const resolved: ResolvedAccount = { apiKey: account.api_key, email: account.email, + accountId, }; cache.set(cacheKey, resolved); diff --git a/src/lib/amplitude-analytics.ts b/src/lib/amplitude-analytics.ts new file mode 100644 index 0000000..de351a6 --- /dev/null +++ b/src/lib/amplitude-analytics.ts @@ -0,0 +1,160 @@ +import { createHash } from 'node:crypto'; +import { + AmplitudeMCPAnalytics, + getCurrentContext, + setIdentity, + setRationale, +} from '@amplitude/mcp-analytics'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import type { FastMCP } from 'fastmcp'; +import type { BrowserlessSession } from '../@types/types.js'; + +export type AmplitudeFactory = ( + apiKey: string, + serverVersion: string, +) => AmplitudeMCPAnalytics; + +const CONNECT_GUARD = Symbol('browserlessAmplitudeConnectGuard'); +const ORIGINAL_CONNECT = Symbol('browserlessAmplitudeOriginalConnect'); +const ADD_TOOL_HOOK = Symbol('browserlessAmplitudeAddToolHook'); +const AMPLITUDE_SHUTDOWN_TIMEOUT_MS = 2_000; + +type HookedServer = Server & { + [CONNECT_GUARD]?: boolean; + [ORIGINAL_CONNECT]?: Server['connect']; +}; + +type HookedFastMcp = FastMCP & { + [ADD_TOOL_HOOK]?: boolean; +}; + +let activeAnalytics: AmplitudeMCPAnalytics | undefined; +let hookInstalled = false; + +const installConnectHook = (analytics: AmplitudeMCPAnalytics): void => { + activeAnalytics = analytics; + if (hookInstalled) return; + + const originalConnect = Server.prototype.connect; + Server.prototype.connect = function (this: HookedServer, transport) { + if (this[CONNECT_GUARD]) { + return this[ORIGINAL_CONNECT]!(transport); + } + + const currentAnalytics = activeAnalytics; + if (!currentAnalytics) { + return originalConnect.call(this, transport); + } + + this[CONNECT_GUARD] = true; + this[ORIGINAL_CONNECT] = originalConnect.bind(this); + try { + currentAnalytics.instrumentServer(this); + } catch (error) { + console.error( + '[browserless-mcp] Amplitude instrumentation failed:', + error, + ); + } + return this.connect(transport); + }; + hookInstalled = true; +}; + +export const initializeAmplitudeAnalytics = ( + apiKey: string | undefined, + serverVersion: string, + factory: AmplitudeFactory = (key, version) => + new AmplitudeMCPAnalytics({ + apiKey: key, + serverName: 'browserless-mcp', + serverVersion: version, + }), +): AmplitudeMCPAnalytics | undefined => { + if (!apiKey) return undefined; + + try { + const analytics = factory(apiKey, serverVersion); + installConnectHook(analytics); + return analytics; + } catch (error) { + console.error('[browserless-mcp] Amplitude initialization failed:', error); + return undefined; + } +}; + +export const instrumentFastMcpTools = ( + server: FastMCP, + analytics: AmplitudeMCPAnalytics | undefined, +): void => { + if (!analytics) return; + + const hookedServer = server as HookedFastMcp; + if (hookedServer[ADD_TOOL_HOOK]) return; + + const originalAddTool = server.addTool.bind(server); + server.addTool = ((tool) => { + const execute = async ( + ...args: Parameters + ): Promise => + // FastMCP and MCP SDK content unions differ although the runtime shape matches. + (await tool.execute(...args)) as CallToolResult; + + return originalAddTool({ + ...tool, + execute: analytics.instrumentTool(execute, { name: tool.name }), + }); + }) as FastMCP['addTool']; + hookedServer[ADD_TOOL_HOOK] = true; +}; + +export const getAmplitudeIdentity = ( + session: BrowserlessSession | undefined, + token: string, +): string => + session?.accountId ?? + `token-${createHash('sha256').update(token).digest('base64url')}`; + +export const setAmplitudeToolContext = ( + session: BrowserlessSession | undefined, + token: string, + prompt: string | undefined, +): void => { + if (!activeAnalytics) return; + + try { + if (!getCurrentContext()) return; + setIdentity({ userId: getAmplitudeIdentity(session, token) }); + if (prompt !== undefined) setRationale(prompt); + } catch (error) { + console.error('[browserless-mcp] Amplitude tool context failed:', error); + } +}; + +export const shutdownAmplitudeAnalytics = async ( + analytics: AmplitudeMCPAnalytics | undefined, +): Promise => { + let timeout: NodeJS.Timeout | undefined; + try { + await Promise.race([ + analytics?.shutdown(), + new Promise((resolve) => { + timeout = setTimeout(resolve, AMPLITUDE_SHUTDOWN_TIMEOUT_MS); + timeout.unref(); + }), + ]); + } catch (error) { + console.error('[browserless-mcp] Amplitude shutdown failed:', error); + } finally { + if (timeout) clearTimeout(timeout); + } +}; + +export const resetAmplitudeAnalyticsForTests = ( + originalConnect: Server['connect'], +): void => { + Server.prototype.connect = originalConnect; + activeAnalytics = undefined; + hookInstalled = false; +}; diff --git a/src/lib/define-tool.ts b/src/lib/define-tool.ts index 27c4022..4124d1d 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -9,6 +9,7 @@ import { } from './utils.js'; import { ResponseCache } from './cache.js'; import { AnalyticsHelper } from './analytics.js'; +import { setAmplitudeToolContext } from './amplitude-analytics.js'; import type { ApiClient, BrowserlessSession, @@ -159,6 +160,8 @@ export function defineTool( } const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; + setAmplitudeToolContext(s, token, prompt); + def.validateUrl?.(params); await reportProgress({ progress: 0, total: 100 }); diff --git a/src/lib/http-auth.ts b/src/lib/http-auth.ts index 0e33c9b..67357b0 100644 --- a/src/lib/http-auth.ts +++ b/src/lib/http-auth.ts @@ -7,6 +7,7 @@ export interface ResolvedBrowserlessAuth { apiUrl: string; attachSessionId?: string; source?: string; + accountId?: string; } export interface AuthInput { @@ -60,12 +61,12 @@ export const resolveBrowserlessAuth = async ( // A JWT is exchanged for the account's Browserless API key via PostgREST. if (isJwt && headerToken) { - const { apiKey } = await resolveApiKey( + const { apiKey, accountId } = await resolveApiKey( config.supabaseUrl, config.supabaseServiceRoleKey, headerToken, ); - return { token: apiKey, apiUrl, attachSessionId, source }; + return { token: apiKey, apiUrl, attachSessionId, source, accountId }; } throw new Error( diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts new file mode 100644 index 0000000..acd91a3 --- /dev/null +++ b/test/lib/amplitude-analytics.spec.ts @@ -0,0 +1,227 @@ +import { expect } from 'chai'; +import { MockAmplitudeMCPAnalytics } from '@amplitude/mcp-analytics/testing'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { FastMCP } from 'fastmcp'; +import sinon from 'sinon'; +import { + getAmplitudeIdentity, + initializeAmplitudeAnalytics, + instrumentFastMcpTools, + resetAmplitudeAnalyticsForTests, + shutdownAmplitudeAnalytics, +} from '../../src/lib/amplitude-analytics.js'; +import { defineTool } from '../../src/lib/define-tool.js'; +import type { McpConfig } from '../../src/@types/types.js'; +import { z } from 'zod'; + +const createDeferred = () => { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +}; + +describe('Amplitude MCP analytics', () => { + const originalConnect = Server.prototype.connect; + + afterEach(() => { + resetAmplitudeAnalyticsForTests(originalConnect); + }); + + it('is disabled without an API key', () => { + let constructed = false; + const analytics = initializeAmplitudeAnalytics(undefined, '1.0.0', () => { + constructed = true; + return new MockAmplitudeMCPAnalytics({ + serverName: 'test', + serverVersion: '1.0.0', + }); + }); + + expect(analytics).to.equal(undefined); + expect(constructed).to.equal(false); + }); + + it('uses account ids or hashed tokens as identity', () => { + const token = 'plain-browserless-token'; + expect( + getAmplitudeIdentity({ token, apiUrl: 'https://example.com' }, token), + ).to.equal('token-Hv-TiXSqe4TOtrWa7FLy8hqVC8ermzWq7wKnBgIHrX8'); + expect( + getAmplitudeIdentity( + { token, apiUrl: 'https://example.com', accountId: 'account-123' }, + token, + ), + ).to.equal('account-123'); + expect( + getAmplitudeIdentity({ token, apiUrl: 'https://example.com' }, token), + ).to.not.equal(token); + }); + + it('instruments a connected SDK server without recursing', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const instrumentServer = sinon.spy(mock, 'instrumentServer'); + const connect = sinon.spy(originalConnect); + Server.prototype.connect = connect; + initializeAmplitudeAnalytics('test-key', '1.0.0', () => mock); + + const server = new Server({ name: 'test', version: '1.0.0' }); + await server.connect({ + start: async () => undefined, + close: async () => undefined, + send: async () => undefined, + }); + + expect(instrumentServer.calledOnceWithExactly(server)).to.equal(true); + expect(connect.calledOnce).to.equal(true); + }); + + it('instruments FastMCP tools once and skips instrumentation when disabled', () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const instrumentTool = sinon.spy(mock, 'instrumentTool'); + const server = new FastMCP({ name: 'test', version: '1.0.0' }); + const addTool = sinon.spy(server, 'addTool'); + const execute = async () => ({ content: [] }); + + instrumentFastMcpTools(server, mock); + server.addTool({ + name: 'test_tool', + parameters: z.object({}), + execute, + }); + + expect(instrumentTool.calledOnce).to.equal(true); + expect(instrumentTool.firstCall.args[1]).to.deep.equal({ + name: 'test_tool', + }); + expect(addTool.firstCall.args[0].execute).to.not.equal(execute); + + const disabledMock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const disabledInstrumentTool = sinon.spy(disabledMock, 'instrumentTool'); + const disabledServer = new FastMCP({ + name: 'disabled-test', + version: '1.0.0', + }); + const disabledAddTool = sinon.spy(disabledServer, 'addTool'); + const disabledExecute = async () => ({ content: [] }); + + instrumentFastMcpTools(disabledServer, undefined); + disabledServer.addTool({ + name: 'disabled_tool', + parameters: z.object({}), + execute: disabledExecute, + }); + + expect(disabledInstrumentTool.notCalled).to.equal(true); + expect(disabledAddTool.firstCall.args[0].execute).to.equal(disabledExecute); + }); + + it('executes a defined tool normally when analytics is disabled', async () => { + const server = new FastMCP({ name: 'disabled-test', version: '1.0.0' }); + const addTool = sinon.spy(server, 'addTool'); + const config = { + browserlessApiUrl: 'https://example.com', + complianceMode: true, + } as McpConfig; + + defineTool(server, config, undefined, { + name: 'disabled_tool', + description: 'Disabled analytics test', + parameters: z.object({}), + run: async () => 'ok', + format: (result) => [{ type: 'text', text: result }], + }); + + const tool = addTool.firstCall.args[0]; + const result = await tool.execute( + {}, + { + reportProgress: async () => undefined, + session: { token: 'test-token', apiUrl: 'https://example.com' }, + sessionId: undefined, + log: { + debug: () => undefined, + error: () => undefined, + info: () => undefined, + warn: () => undefined, + }, + client: { version: { name: 'test-client', version: '1.0.0' } }, + streamContent: async () => undefined, + }, + ); + + expect(result).to.deep.equal({ + content: [{ type: 'text', text: 'ok' }], + }); + }); + + it('awaits successful Amplitude shutdown', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const deferred = createDeferred(); + const shutdown = sinon.stub(mock, 'shutdown') as sinon.SinonStub; + shutdown.returns(deferred.promise); + const result = shutdownAmplitudeAnalytics(mock); + + expect( + await Promise.race([ + result.then(() => 'settled'), + Promise.resolve('pending'), + ]), + ).to.equal('pending'); + deferred.resolve(); + await result; + + expect(shutdown.calledOnce).to.equal(true); + }); + + it('swallows rejected Amplitude shutdown', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const deferred = createDeferred(); + const shutdown = sinon.stub(mock, 'shutdown') as sinon.SinonStub; + shutdown.returns(deferred.promise); + const result = shutdownAmplitudeAnalytics(mock); + + expect( + await Promise.race([ + result.then(() => 'settled'), + Promise.resolve('pending'), + ]), + ).to.equal('pending'); + deferred.reject(new Error('flush failed')); + await result; + + expect(shutdown.calledOnce).to.equal(true); + }); + + it('resolves when Amplitude shutdown never settles', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const shutdown = sinon.stub(mock, 'shutdown') as sinon.SinonStub; + shutdown.returns(new Promise(() => {})); + + const result = await shutdownAmplitudeAnalytics(mock); + + expect(result).to.equal(undefined); + expect(shutdown.calledOnce).to.equal(true); + }); +});