From f11d47fdbaebe1a4ab179915476c6decf2f194bc Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:27:07 +0000 Subject: [PATCH 1/9] feat: add Amplitude MCP analytics Co-Authored-By: Sean Yang --- README.md | 1 + llms-install.md | 3 +- package-lock.json | 75 +++++++++++++++++++-- package.json | 3 + src/@types/types.d.ts | 3 + src/config.ts | 1 + src/index.ts | 31 +++++++++ src/lib/account-resolver.ts | 2 + src/lib/amplitude-analytics.ts | 93 ++++++++++++++++++++++++++ src/lib/define-tool.ts | 98 ++++++++++++++++++---------- src/lib/http-auth.ts | 5 +- test/lib/amplitude-analytics.spec.ts | 67 +++++++++++++++++++ 12 files changed, 340 insertions(+), 42 deletions(-) create mode 100644 src/lib/amplitude-analytics.ts create mode 100644 test/lib/amplitude-analytics.spec.ts 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..bc5a119 100644 --- a/llms-install.md +++ b/llms-install.md @@ -45,7 +45,8 @@ If the user needs a local-only / air-gapped install, use the npm package over st "command": "npx", "args": ["browserless-mcp"], "env": { - "BROWSERLESS_TOKEN": "" + "BROWSERLESS_TOKEN": "", + "AMPLITUDE_API_KEY": "" }, "disabled": false, "autoApprove": [] diff --git a/package-lock.json b/package-lock.json index a9469b4..ba7278e 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", @@ -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..5f8be56 100644 --- a/src/index.ts +++ b/src/index.ts @@ -16,6 +16,11 @@ 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 { + initializeAmplitudeAnalytics, + shutdownAmplitudeAnalytics, +} from './lib/amplitude-analytics.js'; const pkg = JSON.parse( readFileSync( @@ -38,6 +43,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. @@ -143,15 +152,37 @@ 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 = (): void => { + if (amplitudeShutdown) return; + amplitudeShutdown = true; + shutdownAmplitudeAnalytics(amplitudeAnalytics); + }; + process.once('SIGTERM', shutdown); + process.once('SIGINT', shutdown); +} + 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..50f3955 --- /dev/null +++ b/src/lib/amplitude-analytics.ts @@ -0,0 +1,93 @@ +import { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import type { BrowserlessSession } from '../@types/types.js'; +import { djb2 } from './utils.js'; + +type AmplitudeServer = AmplitudeMCPAnalytics & { + instrumentServer: (server: Server) => Server; +}; +export type AmplitudeFactory = ( + apiKey: string, + serverVersion: string, +) => AmplitudeServer; + +const CONNECT_GUARD = Symbol('browserlessAmplitudeConnectGuard'); +const ORIGINAL_CONNECT = Symbol('browserlessAmplitudeOriginalConnect'); + +type HookedServer = Server & { + [CONNECT_GUARD]?: boolean; + [ORIGINAL_CONNECT]?: Server['connect']; +}; + +let activeAnalytics: AmplitudeServer | undefined; +let hookInstalled = false; + +const installConnectHook = (analytics: AmplitudeServer): 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, + }) as AmplitudeServer, +): AmplitudeServer | 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 getAmplitudeAnalytics = (): AmplitudeServer | undefined => + activeAnalytics; + +export const getAmplitudeIdentity = ( + session: BrowserlessSession | undefined, + token: string, +): string => session?.accountId ?? String(djb2(token)); + +export const shutdownAmplitudeAnalytics = ( + analytics: AmplitudeServer | undefined, +): void => { + try { + analytics?.shutdown(); + } catch (error) { + console.error('[browserless-mcp] Amplitude shutdown failed:', error); + } +}; diff --git a/src/lib/define-tool.ts b/src/lib/define-tool.ts index 27c4022..d1673ee 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -1,5 +1,5 @@ import { FastMCP, UserError } from 'fastmcp'; -import type { Content } from 'fastmcp'; +import type { Content, Context } from 'fastmcp'; import { z, type ZodType } from 'zod'; import { createApiClient, ProfileNotFoundError } from './api-client.js'; import { @@ -9,6 +9,10 @@ import { } from './utils.js'; import { ResponseCache } from './cache.js'; import { AnalyticsHelper } from './analytics.js'; +import { + getAmplitudeAnalytics, + getAmplitudeIdentity, +} from './amplitude-analytics.js'; import type { ApiClient, BrowserlessSession, @@ -129,15 +133,18 @@ export function defineTool( ? def.parameters.extend({ _prompt: PROMPT_FIELD }) : def.parameters; - server.addTool({ - name: def.name, - description: def.description, - parameters, - annotations: def.annotations, - execute: async ( - args, - { reportProgress, session, sessionId, log, client: mcpClient }, - ) => { + const amplitude = getAmplitudeAnalytics(); + const execute = async ( + args: Record | P, + { + reportProgress, + session, + sessionId, + log, + client: mcpClient, + }: Context | undefined>, + ) => { + try { // Split the injected `_prompt` off so it never reaches `run`/the API. const { _prompt, ...rest } = (args ?? {}) as Record; const prompt = @@ -159,6 +166,18 @@ export function defineTool( } const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; + if (amplitude) { + try { + amplitude.setIdentity({ userId: getAmplitudeIdentity(s, token) }); + if (prompt !== undefined) amplitude.setRationale(prompt); + } catch (error) { + console.error( + '[browserless-mcp] Amplitude tool context failed:', + error, + ); + } + } + def.validateUrl?.(params); await reportProgress({ progress: 0, total: 100 }); @@ -172,30 +191,19 @@ export function defineTool( def.cache, ); - let result: R; - try { - result = await def.run({ - client, - params, - prompt, - log, - analytics, - mcpSource, - token, - apiUrl, - reportProgress, - sessionId, - attachSessionId: s?.attachSessionId, - }); - } catch (err) { - if (err instanceof ProfileNotFoundError) { - const msg = def.profileNotFoundMessage - ? def.profileNotFoundMessage(err.profile) - : defaultProfileMessage(err.profile); - throw new UserError(msg); - } - throw err; - } + const result = await def.run({ + client, + params, + prompt, + log, + analytics, + mcpSource, + token, + apiUrl, + reportProgress, + sessionId, + attachSessionId: s?.attachSessionId, + }); await reportProgress({ progress: 100, total: 100 }); @@ -209,6 +217,26 @@ export function defineTool( } return { content: def.format(result, params) }; - }, + } catch (err) { + if (err instanceof ProfileNotFoundError) { + const msg = def.profileNotFoundMessage + ? def.profileNotFoundMessage(err.profile) + : defaultProfileMessage(err.profile); + throw new UserError(msg); + } + throw err; + } + }; + + server.addTool({ + name: def.name, + description: def.description, + parameters, + annotations: def.annotations, + execute: amplitude + ? (amplitude.instrumentTool(execute as never, { + name: def.name, + }) as typeof execute) + : execute, }); } 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..b08841e --- /dev/null +++ b/test/lib/amplitude-analytics.spec.ts @@ -0,0 +1,67 @@ +import { expect } from 'chai'; +import { MockAmplitudeMCPAnalytics } from '@amplitude/mcp-analytics/testing'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { + getAmplitudeAnalytics, + getAmplitudeIdentity, + initializeAmplitudeAnalytics, +} from '../../src/lib/amplitude-analytics.js'; +import { djb2 } from '../../src/lib/utils.js'; + +describe('Amplitude MCP analytics', () => { + 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', + }) as never; + }); + + 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(String(djb2(token))); + 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', + }); + initializeAmplitudeAnalytics('test-key', '1.0.0', () => mock as never); + + const server = new Server({ name: 'test', version: '1.0.0' }); + await server.connect({ + start: async () => undefined, + close: async () => undefined, + send: async () => undefined, + }); + + const wrapped = mock.instrumentTool( + async (_args: Record, _extra: unknown) => ({ + content: [{ type: 'text', text: 'ok' }], + }), + { name: 'search' }, + ); + await wrapped({}, { sessionId: 'session-1' } as never); + + expect(getAmplitudeAnalytics()).to.equal(mock); + expect(mock.getEvents('[MCP] Tool Call Response')).to.have.length(1); + }); +}); From f71df7590036eb300fa1cc026dbcf1d88fb14ead Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:32:49 +0000 Subject: [PATCH 2/9] fix: harden Amplitude analytics integration Co-Authored-By: Sean Yang --- src/index.ts | 16 ++-- src/lib/amplitude-analytics.ts | 34 ++++---- src/lib/define-tool.ts | 126 +++++++++++++-------------- src/tools/agent.ts | 12 ++- src/tools/crawl.ts | 15 +++- src/tools/export.ts | 15 +++- src/tools/function.ts | 12 ++- src/tools/map.ts | 15 +++- src/tools/performance.ts | 15 +++- src/tools/profiles.ts | 12 ++- src/tools/register.ts | 24 +++-- src/tools/search.ts | 12 ++- src/tools/smartscraper.ts | 15 +++- test/lib/amplitude-analytics.spec.ts | 30 ++++--- 14 files changed, 237 insertions(+), 116 deletions(-) diff --git a/src/index.ts b/src/index.ts index 5f8be56..3bd81f0 100644 --- a/src/index.ts +++ b/src/index.ts @@ -131,7 +131,7 @@ const server = new FastMCP({ authenticate: hybridAuthenticate, }); -registerSurface(server, config, analytics); +registerSurface(server, config, analytics, amplitudeAnalytics); // 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 // on a directory deploy) from opt-out, and warn on an unrecognized value (typo). @@ -174,13 +174,19 @@ server.on('connect', (event) => { if (amplitudeAnalytics) { let amplitudeShutdown = false; - const shutdown = (): void => { + const shutdown = (exitCode: number): void => { if (amplitudeShutdown) return; amplitudeShutdown = true; - shutdownAmplitudeAnalytics(amplitudeAnalytics); + void (async () => { + try { + await shutdownAmplitudeAnalytics(amplitudeAnalytics); + } finally { + process.exit(exitCode); + } + })(); }; - process.once('SIGTERM', shutdown); - process.once('SIGINT', shutdown); + process.once('SIGTERM', () => shutdown(143)); + process.once('SIGINT', () => shutdown(130)); } server.on('disconnect', (event) => { diff --git a/src/lib/amplitude-analytics.ts b/src/lib/amplitude-analytics.ts index 50f3955..d10c979 100644 --- a/src/lib/amplitude-analytics.ts +++ b/src/lib/amplitude-analytics.ts @@ -3,13 +3,10 @@ import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import type { BrowserlessSession } from '../@types/types.js'; import { djb2 } from './utils.js'; -type AmplitudeServer = AmplitudeMCPAnalytics & { - instrumentServer: (server: Server) => Server; -}; export type AmplitudeFactory = ( apiKey: string, serverVersion: string, -) => AmplitudeServer; +) => AmplitudeMCPAnalytics; const CONNECT_GUARD = Symbol('browserlessAmplitudeConnectGuard'); const ORIGINAL_CONNECT = Symbol('browserlessAmplitudeOriginalConnect'); @@ -19,10 +16,10 @@ type HookedServer = Server & { [ORIGINAL_CONNECT]?: Server['connect']; }; -let activeAnalytics: AmplitudeServer | undefined; +let activeAnalytics: AmplitudeMCPAnalytics | undefined; let hookInstalled = false; -const installConnectHook = (analytics: AmplitudeServer): void => { +const installConnectHook = (analytics: AmplitudeMCPAnalytics): void => { activeAnalytics = analytics; if (hookInstalled) return; @@ -60,8 +57,8 @@ export const initializeAmplitudeAnalytics = ( apiKey: key, serverName: 'browserless-mcp', serverVersion: version, - }) as AmplitudeServer, -): AmplitudeServer | undefined => { + }), +): AmplitudeMCPAnalytics | undefined => { if (!apiKey) return undefined; try { @@ -74,20 +71,25 @@ export const initializeAmplitudeAnalytics = ( } }; -export const getAmplitudeAnalytics = (): AmplitudeServer | undefined => - activeAnalytics; - export const getAmplitudeIdentity = ( session: BrowserlessSession | undefined, token: string, -): string => session?.accountId ?? String(djb2(token)); +): string => session?.accountId ?? `token-${djb2(token)}`; -export const shutdownAmplitudeAnalytics = ( - analytics: AmplitudeServer | undefined, -): void => { +export const shutdownAmplitudeAnalytics = async ( + analytics: AmplitudeMCPAnalytics | undefined, +): Promise => { try { - analytics?.shutdown(); + await analytics?.shutdown(); } catch (error) { console.error('[browserless-mcp] Amplitude shutdown failed:', error); } }; + +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 d1673ee..4280940 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -1,5 +1,6 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content, Context } from 'fastmcp'; +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; import { z, type ZodType } from 'zod'; import { createApiClient, ProfileNotFoundError } from './api-client.js'; import { @@ -9,10 +10,8 @@ import { } from './utils.js'; import { ResponseCache } from './cache.js'; import { AnalyticsHelper } from './analytics.js'; -import { - getAmplitudeAnalytics, - getAmplitudeIdentity, -} from './amplitude-analytics.js'; +import { getAmplitudeIdentity } from './amplitude-analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ApiClient, BrowserlessSession, @@ -125,6 +124,7 @@ export function defineTool( config: McpConfig, analytics: AnalyticsHelper | undefined, def: ToolDefinition, + amplitude?: AmplitudeMCPAnalytics, ): void { // Not on the compliant surface: it's a strict allowlist / privacy gate, so // we don't ask the model to self-report user prompts there. @@ -133,7 +133,6 @@ export function defineTool( ? def.parameters.extend({ _prompt: PROMPT_FIELD }) : def.parameters; - const amplitude = getAmplitudeAnalytics(); const execute = async ( args: Record | P, { @@ -144,54 +143,55 @@ export function defineTool( client: mcpClient, }: Context | undefined>, ) => { - try { - // Split the injected `_prompt` off so it never reaches `run`/the API. - const { _prompt, ...rest } = (args ?? {}) as Record; - const prompt = - typeof _prompt === 'string' ? redactSecrets(_prompt) : undefined; - const params = rest as P; - // Single localized cast — FastMCP types session as Record - // for the unconstrained generic. Tools see the typed session via this helper - // and never cast token/apiUrl themselves. - const s = session as BrowserlessSession | undefined; - const mcpSource = resolveMcpSource(s?.source, mcpClient?.version); - - const token = s?.token ?? config.browserlessToken; - if (!token) { - throw new UserError( - 'No Browserless API token provided. ' + - 'For stdio: set the BROWSERLESS_TOKEN environment variable. ' + - 'For HTTP: pass Authorization: Bearer header.', + // Split the injected `_prompt` off so it never reaches `run`/the API. + const { _prompt, ...rest } = (args ?? {}) as Record; + const prompt = + typeof _prompt === 'string' ? redactSecrets(_prompt) : undefined; + const params = rest as P; + // Single localized cast — FastMCP types session as Record + // for the unconstrained generic. Tools see the typed session via this helper + // and never cast token/apiUrl themselves. + const s = session as BrowserlessSession | undefined; + const mcpSource = resolveMcpSource(s?.source, mcpClient?.version); + + const token = s?.token ?? config.browserlessToken; + if (!token) { + throw new UserError( + 'No Browserless API token provided. ' + + 'For stdio: set the BROWSERLESS_TOKEN environment variable. ' + + 'For HTTP: pass Authorization: Bearer header.', + ); + } + const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; + + if (amplitude) { + try { + amplitude.setIdentity({ userId: getAmplitudeIdentity(s, token) }); + if (prompt !== undefined) amplitude.setRationale(prompt); + } catch (error) { + console.error( + '[browserless-mcp] Amplitude tool context failed:', + error, ); } - const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; - - if (amplitude) { - try { - amplitude.setIdentity({ userId: getAmplitudeIdentity(s, token) }); - if (prompt !== undefined) amplitude.setRationale(prompt); - } catch (error) { - console.error( - '[browserless-mcp] Amplitude tool context failed:', - error, - ); - } - } + } - def.validateUrl?.(params); + def.validateUrl?.(params); - await reportProgress({ progress: 0, total: 100 }); + await reportProgress({ progress: 0, total: 100 }); - const client = createApiClient( - { - ...config, - browserlessToken: token, - browserlessApiUrl: apiUrl, - }, - def.cache, - ); + const client = createApiClient( + { + ...config, + browserlessToken: token, + browserlessApiUrl: apiUrl, + }, + def.cache, + ); - const result = await def.run({ + let result: R; + try { + result = await def.run({ client, params, prompt, @@ -204,19 +204,6 @@ export function defineTool( sessionId, attachSessionId: s?.attachSessionId, }); - - await reportProgress({ progress: 100, total: 100 }); - - if (analytics && def.analyticsProps) { - analytics.fireToolRequest(token, def.name, { - api_url: apiUrl, - ...mcpSource, - ...(prompt ? { _prompt: prompt } : {}), - ...def.analyticsProps(params, result), - }); - } - - return { content: def.format(result, params) }; } catch (err) { if (err instanceof ProfileNotFoundError) { const msg = def.profileNotFoundMessage @@ -226,6 +213,21 @@ export function defineTool( } throw err; } + + await reportProgress({ progress: 100, total: 100 }); + + if (analytics && def.analyticsProps) { + analytics.fireToolRequest(token, def.name, { + api_url: apiUrl, + ...mcpSource, + ...(prompt ? { _prompt: prompt } : {}), + ...def.analyticsProps(params, result), + }); + } + + // FastMCP's content union is narrower than the SDK result type expected + // by Amplitude, although the runtime payload is the same MCP shape. + return { content: def.format(result, params) } as CallToolResult; }; server.addTool({ @@ -234,9 +236,7 @@ export function defineTool( parameters, annotations: def.annotations, execute: amplitude - ? (amplitude.instrumentTool(execute as never, { - name: def.name, - }) as typeof execute) + ? amplitude.instrumentTool(execute, { name: def.name }) : execute, }); } diff --git a/src/tools/agent.ts b/src/tools/agent.ts index abb58b9..30d98ec 100644 --- a/src/tools/agent.ts +++ b/src/tools/agent.ts @@ -25,7 +25,9 @@ import type { } from '../@types/types.js'; import { classifyAgentError } from '../lib/error-classifier.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import { defineTool } from '../lib/define-tool.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { defineTool as defineToolBase } from '../lib/define-tool.js'; import { markFired, renderSkill, @@ -425,7 +427,15 @@ export function registerAgentTools( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); // Compliant: drop circumvention/autologin recipes from the enum + omit the diff --git a/src/tools/crawl.ts b/src/tools/crawl.ts index 58bf552..18de37e 100644 --- a/src/tools/crawl.ts +++ b/src/tools/crawl.ts @@ -1,9 +1,14 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { + defineTool as defineToolBase, + validateHttpUrl, +} from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { CrawlPageResult, CrawlParams, @@ -221,7 +226,15 @@ export function registerCrawlTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_crawl', description: diff --git a/src/tools/export.ts b/src/tools/export.ts index a2a0872..4180ff8 100644 --- a/src/tools/export.ts +++ b/src/tools/export.ts @@ -1,10 +1,15 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { + defineTool as defineToolBase, + validateHttpUrl, +} from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { isCompliant, COMPLIANT_EXPORT_DESCRIPTION } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ExportParams, GenericApiResult, @@ -76,7 +81,15 @@ export function registerExportTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool(server, config, analytics, { diff --git a/src/tools/function.ts b/src/tools/function.ts index d09f6a2..59caf5f 100644 --- a/src/tools/function.ts +++ b/src/tools/function.ts @@ -1,9 +1,11 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { defineTool as defineToolBase } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { FunctionParams, GenericApiResult, @@ -116,7 +118,15 @@ export function registerFunctionTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_function', description: diff --git a/src/tools/map.ts b/src/tools/map.ts index 398ed86..1c00cc2 100644 --- a/src/tools/map.ts +++ b/src/tools/map.ts @@ -1,8 +1,13 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { + defineTool as defineToolBase, + validateHttpUrl, +} from '../lib/define-tool.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { MapParams, MapResponse, McpConfig } from '../@types/types.js'; export const SitemapModeSchema = z.enum(['include', 'skip', 'only']); @@ -48,7 +53,15 @@ export function registerMapTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_map', description: diff --git a/src/tools/performance.ts b/src/tools/performance.ts index 81f9bdf..8e25fd6 100644 --- a/src/tools/performance.ts +++ b/src/tools/performance.ts @@ -1,10 +1,15 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { + defineTool as defineToolBase, + validateHttpUrl, +} from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { isCompliant } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, PerformanceParams, @@ -57,7 +62,15 @@ export function registerPerformanceTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool( server, diff --git a/src/tools/profiles.ts b/src/tools/profiles.ts index d91f113..ac1b5cd 100644 --- a/src/tools/profiles.ts +++ b/src/tools/profiles.ts @@ -1,7 +1,9 @@ import { FastMCP } from 'fastmcp'; import { z } from 'zod'; -import { defineTool } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { defineTool as defineToolBase } from '../lib/define-tool.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ListProfilesRequest, McpConfig, @@ -28,7 +30,15 @@ export function registerProfilesTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_profiles', description: diff --git a/src/tools/register.ts b/src/tools/register.ts index d7816ad..9115752 100644 --- a/src/tools/register.ts +++ b/src/tools/register.ts @@ -1,6 +1,7 @@ import { FastMCP } from 'fastmcp'; import type { McpConfig } from '../@types/types.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import { registerSmartScraperTool } from './smartscraper.js'; import { registerExportTool } from './export.js'; import { registerAgentTools } from './agent.js'; @@ -25,6 +26,7 @@ export function registerSurface( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { const registrations: ReadonlyArray<{ surface: Surface; @@ -32,20 +34,21 @@ export function registerSurface( }> = [ { surface: 'both', - register: () => registerExportTool(server, config, analytics), + register: () => registerExportTool(server, config, analytics, amplitude), }, // agent + skill tools { surface: 'both', - register: () => registerAgentTools(server, config, analytics), + register: () => registerAgentTools(server, config, analytics, amplitude), }, { surface: 'both', - register: () => registerSearchTool(server, config, analytics), + register: () => registerSearchTool(server, config, analytics, amplitude), }, { surface: 'both', - register: () => registerPerformanceTool(server, config, analytics), + register: () => + registerPerformanceTool(server, config, analytics, amplitude), }, // Generic service status — safe on both (it reports the active surface). { surface: 'both', register: () => registerStatusResource(server, config) }, @@ -53,24 +56,27 @@ export function registerSurface( // (they document proxy/captcha and steer the model to smartscraper). { surface: 'full', - register: () => registerSmartScraperTool(server, config, analytics), + register: () => + registerSmartScraperTool(server, config, analytics, amplitude), }, { surface: 'full', - register: () => registerFunctionTool(server, config, analytics), + register: () => + registerFunctionTool(server, config, analytics, amplitude), }, { surface: 'full', - register: () => registerMapTool(server, config, analytics), + register: () => registerMapTool(server, config, analytics, amplitude), }, { surface: 'full', - register: () => registerCrawlTool(server, config, analytics), + register: () => registerCrawlTool(server, config, analytics, amplitude), }, // Full-only: the compliant surface has no profile capability (agent rejects `profile`). { surface: 'full', - register: () => registerProfilesTool(server, config, analytics), + register: () => + registerProfilesTool(server, config, analytics, amplitude), }, { surface: 'full', diff --git a/src/tools/search.ts b/src/tools/search.ts index 97840e2..0f73ccb 100644 --- a/src/tools/search.ts +++ b/src/tools/search.ts @@ -1,9 +1,11 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { defineTool as defineToolBase } from '../lib/define-tool.js'; import { isCompliant, COMPLIANT_SEARCH_DESCRIPTION } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, SearchParams, @@ -99,7 +101,15 @@ export function registerSearchTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool(server, config, analytics, { diff --git a/src/tools/smartscraper.ts b/src/tools/smartscraper.ts index 998e54b..9ed18b7 100644 --- a/src/tools/smartscraper.ts +++ b/src/tools/smartscraper.ts @@ -1,10 +1,15 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; +import type { ToolDefinition } from '../lib/define-tool.js'; +import { + defineTool as defineToolBase, + validateHttpUrl, +} from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { ResponseCache } from '../lib/cache.js'; import { AnalyticsHelper } from '../lib/analytics.js'; +import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, SmartScrapeResult, @@ -60,7 +65,15 @@ export function registerSmartScraperTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, + amplitude?: AmplitudeMCPAnalytics, ): void { + const defineTool = ( + toolServer: FastMCP, + toolConfig: McpConfig, + toolAnalytics: AnalyticsHelper | undefined, + def: ToolDefinition, + ): void => + defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const cache = new ResponseCache(config.cacheTtlMs); defineTool(server, config, analytics, { diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index b08841e..59dd8ed 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -1,14 +1,21 @@ import { expect } from 'chai'; import { MockAmplitudeMCPAnalytics } from '@amplitude/mcp-analytics/testing'; import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import sinon from 'sinon'; import { - getAmplitudeAnalytics, getAmplitudeIdentity, initializeAmplitudeAnalytics, + resetAmplitudeAnalyticsForTests, } from '../../src/lib/amplitude-analytics.js'; import { djb2 } from '../../src/lib/utils.js'; 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', () => { @@ -16,7 +23,7 @@ describe('Amplitude MCP analytics', () => { return new MockAmplitudeMCPAnalytics({ serverName: 'test', serverVersion: '1.0.0', - }) as never; + }); }); expect(analytics).to.equal(undefined); @@ -27,7 +34,7 @@ describe('Amplitude MCP analytics', () => { const token = 'plain-browserless-token'; expect( getAmplitudeIdentity({ token, apiUrl: 'https://example.com' }, token), - ).to.equal(String(djb2(token))); + ).to.equal(`token-${djb2(token)}`); expect( getAmplitudeIdentity( { token, apiUrl: 'https://example.com', accountId: 'account-123' }, @@ -44,7 +51,10 @@ describe('Amplitude MCP analytics', () => { serverName: 'browserless-mcp', serverVersion: '1.0.0', }); - initializeAmplitudeAnalytics('test-key', '1.0.0', () => mock as never); + 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({ @@ -53,15 +63,7 @@ describe('Amplitude MCP analytics', () => { send: async () => undefined, }); - const wrapped = mock.instrumentTool( - async (_args: Record, _extra: unknown) => ({ - content: [{ type: 'text', text: 'ok' }], - }), - { name: 'search' }, - ); - await wrapped({}, { sessionId: 'session-1' } as never); - - expect(getAmplitudeAnalytics()).to.equal(mock); - expect(mock.getEvents('[MCP] Tool Call Response')).to.have.length(1); + expect(instrumentServer.calledOnceWithExactly(server)).to.equal(true); + expect(connect.calledOnce).to.equal(true); }); }); From b22a11212a7ac6d4f72cdbb2e8845f23845e3be7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:36:02 +0000 Subject: [PATCH 3/9] refactor: instrument FastMCP tools centrally Co-Authored-By: Sean Yang --- src/index.ts | 4 ++- src/lib/amplitude-analytics.ts | 32 ++++++++++++++++++ src/lib/define-tool.ts | 17 +++++----- src/tools/agent.ts | 12 +------ src/tools/crawl.ts | 15 +-------- src/tools/export.ts | 15 +-------- src/tools/function.ts | 12 +------ src/tools/map.ts | 15 +-------- src/tools/performance.ts | 15 +-------- src/tools/profiles.ts | 12 +------ src/tools/register.ts | 24 +++++--------- src/tools/search.ts | 12 +------ src/tools/smartscraper.ts | 15 +-------- test/lib/amplitude-analytics.spec.ts | 49 ++++++++++++++++++++++++++++ 14 files changed, 111 insertions(+), 138 deletions(-) diff --git a/src/index.ts b/src/index.ts index 3bd81f0..e7c62ff 100644 --- a/src/index.ts +++ b/src/index.ts @@ -18,6 +18,7 @@ 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'; @@ -131,7 +132,8 @@ const server = new FastMCP({ authenticate: hybridAuthenticate, }); -registerSurface(server, config, analytics, amplitudeAnalytics); +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 // on a directory deploy) from opt-out, and warn on an unrecognized value (typo). diff --git a/src/lib/amplitude-analytics.ts b/src/lib/amplitude-analytics.ts index d10c979..917a343 100644 --- a/src/lib/amplitude-analytics.ts +++ b/src/lib/amplitude-analytics.ts @@ -1,5 +1,7 @@ import { AmplitudeMCPAnalytics } 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'; import { djb2 } from './utils.js'; @@ -10,12 +12,17 @@ export type AmplitudeFactory = ( const CONNECT_GUARD = Symbol('browserlessAmplitudeConnectGuard'); const ORIGINAL_CONNECT = Symbol('browserlessAmplitudeOriginalConnect'); +const ADD_TOOL_HOOK = Symbol('browserlessAmplitudeAddToolHook'); type HookedServer = Server & { [CONNECT_GUARD]?: boolean; [ORIGINAL_CONNECT]?: Server['connect']; }; +type HookedFastMcp = FastMCP & { + [ADD_TOOL_HOOK]?: boolean; +}; + let activeAnalytics: AmplitudeMCPAnalytics | undefined; let hookInstalled = false; @@ -71,6 +78,31 @@ export const initializeAmplitudeAnalytics = ( } }; +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, diff --git a/src/lib/define-tool.ts b/src/lib/define-tool.ts index 4280940..7f2e76a 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -1,6 +1,11 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content, Context } from 'fastmcp'; import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import { + getCurrentContext, + setIdentity, + setRationale, +} from '@amplitude/mcp-analytics'; import { z, type ZodType } from 'zod'; import { createApiClient, ProfileNotFoundError } from './api-client.js'; import { @@ -11,7 +16,6 @@ import { import { ResponseCache } from './cache.js'; import { AnalyticsHelper } from './analytics.js'; import { getAmplitudeIdentity } from './amplitude-analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ApiClient, BrowserlessSession, @@ -124,7 +128,6 @@ export function defineTool( config: McpConfig, analytics: AnalyticsHelper | undefined, def: ToolDefinition, - amplitude?: AmplitudeMCPAnalytics, ): void { // Not on the compliant surface: it's a strict allowlist / privacy gate, so // we don't ask the model to self-report user prompts there. @@ -164,10 +167,10 @@ export function defineTool( } const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; - if (amplitude) { + if (getCurrentContext()) { try { - amplitude.setIdentity({ userId: getAmplitudeIdentity(s, token) }); - if (prompt !== undefined) amplitude.setRationale(prompt); + setIdentity({ userId: getAmplitudeIdentity(s, token) }); + if (prompt !== undefined) setRationale(prompt); } catch (error) { console.error( '[browserless-mcp] Amplitude tool context failed:', @@ -235,8 +238,6 @@ export function defineTool( description: def.description, parameters, annotations: def.annotations, - execute: amplitude - ? amplitude.instrumentTool(execute, { name: def.name }) - : execute, + execute, }); } diff --git a/src/tools/agent.ts b/src/tools/agent.ts index 30d98ec..abb58b9 100644 --- a/src/tools/agent.ts +++ b/src/tools/agent.ts @@ -25,9 +25,7 @@ import type { } from '../@types/types.js'; import { classifyAgentError } from '../lib/error-classifier.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { defineTool as defineToolBase } from '../lib/define-tool.js'; +import { defineTool } from '../lib/define-tool.js'; import { markFired, renderSkill, @@ -427,15 +425,7 @@ export function registerAgentTools( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); // Compliant: drop circumvention/autologin recipes from the enum + omit the diff --git a/src/tools/crawl.ts b/src/tools/crawl.ts index 18de37e..58bf552 100644 --- a/src/tools/crawl.ts +++ b/src/tools/crawl.ts @@ -1,14 +1,9 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { - defineTool as defineToolBase, - validateHttpUrl, -} from '../lib/define-tool.js'; +import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { CrawlPageResult, CrawlParams, @@ -226,15 +221,7 @@ export function registerCrawlTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_crawl', description: diff --git a/src/tools/export.ts b/src/tools/export.ts index 4180ff8..a2a0872 100644 --- a/src/tools/export.ts +++ b/src/tools/export.ts @@ -1,15 +1,10 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { - defineTool as defineToolBase, - validateHttpUrl, -} from '../lib/define-tool.js'; +import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { isCompliant, COMPLIANT_EXPORT_DESCRIPTION } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ExportParams, GenericApiResult, @@ -81,15 +76,7 @@ export function registerExportTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool(server, config, analytics, { diff --git a/src/tools/function.ts b/src/tools/function.ts index 59caf5f..d09f6a2 100644 --- a/src/tools/function.ts +++ b/src/tools/function.ts @@ -1,11 +1,9 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { defineTool as defineToolBase } from '../lib/define-tool.js'; +import { defineTool } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { FunctionParams, GenericApiResult, @@ -118,15 +116,7 @@ export function registerFunctionTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_function', description: diff --git a/src/tools/map.ts b/src/tools/map.ts index 1c00cc2..398ed86 100644 --- a/src/tools/map.ts +++ b/src/tools/map.ts @@ -1,13 +1,8 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { - defineTool as defineToolBase, - validateHttpUrl, -} from '../lib/define-tool.js'; +import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { MapParams, MapResponse, McpConfig } from '../@types/types.js'; export const SitemapModeSchema = z.enum(['include', 'skip', 'only']); @@ -53,15 +48,7 @@ export function registerMapTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_map', description: diff --git a/src/tools/performance.ts b/src/tools/performance.ts index 8e25fd6..81f9bdf 100644 --- a/src/tools/performance.ts +++ b/src/tools/performance.ts @@ -1,15 +1,10 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { - defineTool as defineToolBase, - validateHttpUrl, -} from '../lib/define-tool.js'; +import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { isCompliant } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, PerformanceParams, @@ -62,15 +57,7 @@ export function registerPerformanceTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool( server, diff --git a/src/tools/profiles.ts b/src/tools/profiles.ts index ac1b5cd..d91f113 100644 --- a/src/tools/profiles.ts +++ b/src/tools/profiles.ts @@ -1,9 +1,7 @@ import { FastMCP } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { defineTool as defineToolBase } from '../lib/define-tool.js'; +import { defineTool } from '../lib/define-tool.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { ListProfilesRequest, McpConfig, @@ -30,15 +28,7 @@ export function registerProfilesTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); defineTool(server, config, analytics, { name: 'browserless_profiles', description: diff --git a/src/tools/register.ts b/src/tools/register.ts index 9115752..d7816ad 100644 --- a/src/tools/register.ts +++ b/src/tools/register.ts @@ -1,7 +1,6 @@ import { FastMCP } from 'fastmcp'; import type { McpConfig } from '../@types/types.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import { registerSmartScraperTool } from './smartscraper.js'; import { registerExportTool } from './export.js'; import { registerAgentTools } from './agent.js'; @@ -26,7 +25,6 @@ export function registerSurface( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { const registrations: ReadonlyArray<{ surface: Surface; @@ -34,21 +32,20 @@ export function registerSurface( }> = [ { surface: 'both', - register: () => registerExportTool(server, config, analytics, amplitude), + register: () => registerExportTool(server, config, analytics), }, // agent + skill tools { surface: 'both', - register: () => registerAgentTools(server, config, analytics, amplitude), + register: () => registerAgentTools(server, config, analytics), }, { surface: 'both', - register: () => registerSearchTool(server, config, analytics, amplitude), + register: () => registerSearchTool(server, config, analytics), }, { surface: 'both', - register: () => - registerPerformanceTool(server, config, analytics, amplitude), + register: () => registerPerformanceTool(server, config, analytics), }, // Generic service status — safe on both (it reports the active surface). { surface: 'both', register: () => registerStatusResource(server, config) }, @@ -56,27 +53,24 @@ export function registerSurface( // (they document proxy/captcha and steer the model to smartscraper). { surface: 'full', - register: () => - registerSmartScraperTool(server, config, analytics, amplitude), + register: () => registerSmartScraperTool(server, config, analytics), }, { surface: 'full', - register: () => - registerFunctionTool(server, config, analytics, amplitude), + register: () => registerFunctionTool(server, config, analytics), }, { surface: 'full', - register: () => registerMapTool(server, config, analytics, amplitude), + register: () => registerMapTool(server, config, analytics), }, { surface: 'full', - register: () => registerCrawlTool(server, config, analytics, amplitude), + register: () => registerCrawlTool(server, config, analytics), }, // Full-only: the compliant surface has no profile capability (agent rejects `profile`). { surface: 'full', - register: () => - registerProfilesTool(server, config, analytics, amplitude), + register: () => registerProfilesTool(server, config, analytics), }, { surface: 'full', diff --git a/src/tools/search.ts b/src/tools/search.ts index 0f73ccb..97840e2 100644 --- a/src/tools/search.ts +++ b/src/tools/search.ts @@ -1,11 +1,9 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { defineTool as defineToolBase } from '../lib/define-tool.js'; +import { defineTool } from '../lib/define-tool.js'; import { isCompliant, COMPLIANT_SEARCH_DESCRIPTION } from './compliance.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, SearchParams, @@ -101,15 +99,7 @@ export function registerSearchTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const compliant = isCompliant(config); defineTool(server, config, analytics, { diff --git a/src/tools/smartscraper.ts b/src/tools/smartscraper.ts index 9ed18b7..998e54b 100644 --- a/src/tools/smartscraper.ts +++ b/src/tools/smartscraper.ts @@ -1,15 +1,10 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; import { z } from 'zod'; -import type { ToolDefinition } from '../lib/define-tool.js'; -import { - defineTool as defineToolBase, - validateHttpUrl, -} from '../lib/define-tool.js'; +import { defineTool, validateHttpUrl } from '../lib/define-tool.js'; import { profileField } from './schemas.js'; import { ResponseCache } from '../lib/cache.js'; import { AnalyticsHelper } from '../lib/analytics.js'; -import type { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; import type { McpConfig, SmartScrapeResult, @@ -65,15 +60,7 @@ export function registerSmartScraperTool( server: FastMCP, config: McpConfig, analytics?: AnalyticsHelper, - amplitude?: AmplitudeMCPAnalytics, ): void { - const defineTool = ( - toolServer: FastMCP, - toolConfig: McpConfig, - toolAnalytics: AnalyticsHelper | undefined, - def: ToolDefinition, - ): void => - defineToolBase(toolServer, toolConfig, toolAnalytics, def, amplitude); const cache = new ResponseCache(config.cacheTtlMs); defineTool(server, config, analytics, { diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index 59dd8ed..e7ee742 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -1,13 +1,16 @@ 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, } from '../../src/lib/amplitude-analytics.js'; import { djb2 } from '../../src/lib/utils.js'; +import { z } from 'zod'; describe('Amplitude MCP analytics', () => { const originalConnect = Server.prototype.connect; @@ -66,4 +69,50 @@ describe('Amplitude MCP analytics', () => { 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); + }); }); From ef8558e1eff0efa17604e99abf1841f6cccdc232 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:38:12 +0000 Subject: [PATCH 4/9] cleanup: restore inline tool execution Co-Authored-By: Sean Yang --- src/lib/define-tool.ts | 179 +++++++++++++++++++---------------------- 1 file changed, 84 insertions(+), 95 deletions(-) diff --git a/src/lib/define-tool.ts b/src/lib/define-tool.ts index 7f2e76a..8721572 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -1,6 +1,5 @@ import { FastMCP, UserError } from 'fastmcp'; -import type { Content, Context } from 'fastmcp'; -import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import type { Content } from 'fastmcp'; import { getCurrentContext, setIdentity, @@ -136,108 +135,98 @@ export function defineTool( ? def.parameters.extend({ _prompt: PROMPT_FIELD }) : def.parameters; - const execute = async ( - args: Record | P, - { - reportProgress, - session, - sessionId, - log, - client: mcpClient, - }: Context | undefined>, - ) => { - // Split the injected `_prompt` off so it never reaches `run`/the API. - const { _prompt, ...rest } = (args ?? {}) as Record; - const prompt = - typeof _prompt === 'string' ? redactSecrets(_prompt) : undefined; - const params = rest as P; - // Single localized cast — FastMCP types session as Record - // for the unconstrained generic. Tools see the typed session via this helper - // and never cast token/apiUrl themselves. - const s = session as BrowserlessSession | undefined; - const mcpSource = resolveMcpSource(s?.source, mcpClient?.version); - - const token = s?.token ?? config.browserlessToken; - if (!token) { - throw new UserError( - 'No Browserless API token provided. ' + - 'For stdio: set the BROWSERLESS_TOKEN environment variable. ' + - 'For HTTP: pass Authorization: Bearer header.', - ); - } - const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; - - if (getCurrentContext()) { - try { - setIdentity({ userId: getAmplitudeIdentity(s, token) }); - if (prompt !== undefined) setRationale(prompt); - } catch (error) { - console.error( - '[browserless-mcp] Amplitude tool context failed:', - error, + server.addTool({ + name: def.name, + description: def.description, + parameters, + annotations: def.annotations, + execute: async ( + args, + { reportProgress, session, sessionId, log, client: mcpClient }, + ) => { + // Split the injected `_prompt` off so it never reaches `run`/the API. + const { _prompt, ...rest } = (args ?? {}) as Record; + const prompt = + typeof _prompt === 'string' ? redactSecrets(_prompt) : undefined; + const params = rest as P; + // Single localized cast — FastMCP types session as Record + // for the unconstrained generic. Tools see the typed session via this helper + // and never cast token/apiUrl themselves. + const s = session as BrowserlessSession | undefined; + const mcpSource = resolveMcpSource(s?.source, mcpClient?.version); + + const token = s?.token ?? config.browserlessToken; + if (!token) { + throw new UserError( + 'No Browserless API token provided. ' + + 'For stdio: set the BROWSERLESS_TOKEN environment variable. ' + + 'For HTTP: pass Authorization: Bearer header.', ); } - } + const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; + + if (getCurrentContext()) { + try { + setIdentity({ userId: getAmplitudeIdentity(s, token) }); + if (prompt !== undefined) setRationale(prompt); + } catch (error) { + console.error( + '[browserless-mcp] Amplitude tool context failed:', + error, + ); + } + } - def.validateUrl?.(params); + def.validateUrl?.(params); - await reportProgress({ progress: 0, total: 100 }); + await reportProgress({ progress: 0, total: 100 }); - const client = createApiClient( - { - ...config, - browserlessToken: token, - browserlessApiUrl: apiUrl, - }, - def.cache, - ); + const client = createApiClient( + { + ...config, + browserlessToken: token, + browserlessApiUrl: apiUrl, + }, + def.cache, + ); - let result: R; - try { - result = await def.run({ - client, - params, - prompt, - log, - analytics, - mcpSource, - token, - apiUrl, - reportProgress, - sessionId, - attachSessionId: s?.attachSessionId, - }); - } catch (err) { - if (err instanceof ProfileNotFoundError) { - const msg = def.profileNotFoundMessage - ? def.profileNotFoundMessage(err.profile) - : defaultProfileMessage(err.profile); - throw new UserError(msg); + let result: R; + try { + result = await def.run({ + client, + params, + prompt, + log, + analytics, + mcpSource, + token, + apiUrl, + reportProgress, + sessionId, + attachSessionId: s?.attachSessionId, + }); + } catch (err) { + if (err instanceof ProfileNotFoundError) { + const msg = def.profileNotFoundMessage + ? def.profileNotFoundMessage(err.profile) + : defaultProfileMessage(err.profile); + throw new UserError(msg); + } + throw err; } - throw err; - } - - await reportProgress({ progress: 100, total: 100 }); - if (analytics && def.analyticsProps) { - analytics.fireToolRequest(token, def.name, { - api_url: apiUrl, - ...mcpSource, - ...(prompt ? { _prompt: prompt } : {}), - ...def.analyticsProps(params, result), - }); - } + await reportProgress({ progress: 100, total: 100 }); - // FastMCP's content union is narrower than the SDK result type expected - // by Amplitude, although the runtime payload is the same MCP shape. - return { content: def.format(result, params) } as CallToolResult; - }; + if (analytics && def.analyticsProps) { + analytics.fireToolRequest(token, def.name, { + api_url: apiUrl, + ...mcpSource, + ...(prompt ? { _prompt: prompt } : {}), + ...def.analyticsProps(params, result), + }); + } - server.addTool({ - name: def.name, - description: def.description, - parameters, - annotations: def.annotations, - execute, + return { content: def.format(result, params) }; + }, }); } From 7cb9ac21666ee0f74f891a0733856834d8e20b14 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:41:26 +0000 Subject: [PATCH 5/9] fix: update vulnerable brace expansion Co-Authored-By: Sean Yang --- package-lock.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/package-lock.json b/package-lock.json index ba7278e..257a254 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1468,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": { From 25e5f941fbcf842917450372f492ee63be7a5539 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:52:05 +0000 Subject: [PATCH 6/9] fix: hash anonymous identities cryptographically Co-Authored-By: Sean Yang --- llms-install.md | 5 +++-- src/lib/amplitude-analytics.ts | 6 ++++-- test/lib/amplitude-analytics.spec.ts | 6 ++++-- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/llms-install.md b/llms-install.md index bc5a119..b2493ef 100644 --- a/llms-install.md +++ b/llms-install.md @@ -45,8 +45,7 @@ If the user needs a local-only / air-gapped install, use the npm package over st "command": "npx", "args": ["browserless-mcp"], "env": { - "BROWSERLESS_TOKEN": "", - "AMPLITUDE_API_KEY": "" + "BROWSERLESS_TOKEN": "" }, "disabled": false, "autoApprove": [] @@ -55,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/src/lib/amplitude-analytics.ts b/src/lib/amplitude-analytics.ts index 917a343..b8e1b0e 100644 --- a/src/lib/amplitude-analytics.ts +++ b/src/lib/amplitude-analytics.ts @@ -1,9 +1,9 @@ +import { createHash } from 'node:crypto'; import { AmplitudeMCPAnalytics } 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'; -import { djb2 } from './utils.js'; export type AmplitudeFactory = ( apiKey: string, @@ -106,7 +106,9 @@ export const instrumentFastMcpTools = ( export const getAmplitudeIdentity = ( session: BrowserlessSession | undefined, token: string, -): string => session?.accountId ?? `token-${djb2(token)}`; +): string => + session?.accountId ?? + `token-${createHash('sha256').update(token).digest('base64url')}`; export const shutdownAmplitudeAnalytics = async ( analytics: AmplitudeMCPAnalytics | undefined, diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index e7ee742..faf9f9d 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -1,4 +1,5 @@ import { expect } from 'chai'; +import { createHash } from 'node:crypto'; import { MockAmplitudeMCPAnalytics } from '@amplitude/mcp-analytics/testing'; import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { FastMCP } from 'fastmcp'; @@ -9,7 +10,6 @@ import { instrumentFastMcpTools, resetAmplitudeAnalyticsForTests, } from '../../src/lib/amplitude-analytics.js'; -import { djb2 } from '../../src/lib/utils.js'; import { z } from 'zod'; describe('Amplitude MCP analytics', () => { @@ -37,7 +37,9 @@ describe('Amplitude MCP analytics', () => { const token = 'plain-browserless-token'; expect( getAmplitudeIdentity({ token, apiUrl: 'https://example.com' }, token), - ).to.equal(`token-${djb2(token)}`); + ).to.equal( + `token-${createHash('sha256').update(token).digest('base64url')}`, + ); expect( getAmplitudeIdentity( { token, apiUrl: 'https://example.com', accountId: 'account-123' }, From cd75e20b69951235920fbb446de8a745407d59db Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:56:53 +0000 Subject: [PATCH 7/9] test: cover shutdown and fixed identity hash Co-Authored-By: Sean Yang --- test/lib/amplitude-analytics.spec.ts | 32 ++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index faf9f9d..089c475 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -1,5 +1,4 @@ import { expect } from 'chai'; -import { createHash } from 'node:crypto'; import { MockAmplitudeMCPAnalytics } from '@amplitude/mcp-analytics/testing'; import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { FastMCP } from 'fastmcp'; @@ -9,6 +8,7 @@ import { initializeAmplitudeAnalytics, instrumentFastMcpTools, resetAmplitudeAnalyticsForTests, + shutdownAmplitudeAnalytics, } from '../../src/lib/amplitude-analytics.js'; import { z } from 'zod'; @@ -37,9 +37,7 @@ describe('Amplitude MCP analytics', () => { const token = 'plain-browserless-token'; expect( getAmplitudeIdentity({ token, apiUrl: 'https://example.com' }, token), - ).to.equal( - `token-${createHash('sha256').update(token).digest('base64url')}`, - ); + ).to.equal('token-Hv-TiXSqe4TOtrWa7FLy8hqVC8ermzWq7wKnBgIHrX8'); expect( getAmplitudeIdentity( { token, apiUrl: 'https://example.com', accountId: 'account-123' }, @@ -117,4 +115,30 @@ describe('Amplitude MCP analytics', () => { expect(disabledInstrumentTool.notCalled).to.equal(true); expect(disabledAddTool.firstCall.args[0].execute).to.equal(disabledExecute); }); + + it('awaits successful Amplitude shutdown', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const shutdown = sinon.stub(mock, 'shutdown').resolves(); + + await shutdownAmplitudeAnalytics(mock); + + expect(shutdown.calledOnce).to.equal(true); + }); + + it('swallows rejected Amplitude shutdown', async () => { + const mock = new MockAmplitudeMCPAnalytics({ + serverName: 'browserless-mcp', + serverVersion: '1.0.0', + }); + const shutdown = sinon + .stub(mock, 'shutdown') + .rejects(new Error('flush failed')); + + await shutdownAmplitudeAnalytics(mock); + + expect(shutdown.calledOnce).to.equal(true); + }); }); From 189653c728565f83e475b03606c6827db3809d6d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:01:04 +0000 Subject: [PATCH 8/9] test: verify shutdown waits for completion Co-Authored-By: Sean Yang --- test/lib/amplitude-analytics.spec.ts | 40 +++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index 089c475..c383ac1 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -12,6 +12,16 @@ import { } from '../../src/lib/amplitude-analytics.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; @@ -121,9 +131,19 @@ describe('Amplitude MCP analytics', () => { serverName: 'browserless-mcp', serverVersion: '1.0.0', }); - const shutdown = sinon.stub(mock, 'shutdown').resolves(); + const deferred = createDeferred(); + const shutdown = sinon.stub(mock, 'shutdown') as sinon.SinonStub; + shutdown.returns(deferred.promise); + const result = shutdownAmplitudeAnalytics(mock); - await 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); }); @@ -133,11 +153,19 @@ describe('Amplitude MCP analytics', () => { serverName: 'browserless-mcp', serverVersion: '1.0.0', }); - const shutdown = sinon - .stub(mock, 'shutdown') - .rejects(new Error('flush failed')); + const deferred = createDeferred(); + const shutdown = sinon.stub(mock, 'shutdown') as sinon.SinonStub; + shutdown.returns(deferred.promise); + const result = shutdownAmplitudeAnalytics(mock); - await 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); }); From 86206d0cd83fb7c3699ff9cc7e69de68c60ad657 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:45:04 +0000 Subject: [PATCH 9/9] fix: harden Amplitude shutdown and tool context Co-Authored-By: Sean Yang --- src/lib/amplitude-analytics.ts | 35 +++++++++++++++++- src/lib/define-tool.ts | 19 +--------- test/lib/amplitude-analytics.spec.ts | 55 ++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 19 deletions(-) diff --git a/src/lib/amplitude-analytics.ts b/src/lib/amplitude-analytics.ts index b8e1b0e..de351a6 100644 --- a/src/lib/amplitude-analytics.ts +++ b/src/lib/amplitude-analytics.ts @@ -1,5 +1,10 @@ import { createHash } from 'node:crypto'; -import { AmplitudeMCPAnalytics } from '@amplitude/mcp-analytics'; +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'; @@ -13,6 +18,7 @@ export type AmplitudeFactory = ( 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; @@ -110,13 +116,38 @@ export const getAmplitudeIdentity = ( 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 analytics?.shutdown(); + 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); } }; diff --git a/src/lib/define-tool.ts b/src/lib/define-tool.ts index 8721572..4124d1d 100644 --- a/src/lib/define-tool.ts +++ b/src/lib/define-tool.ts @@ -1,10 +1,5 @@ import { FastMCP, UserError } from 'fastmcp'; import type { Content } from 'fastmcp'; -import { - getCurrentContext, - setIdentity, - setRationale, -} from '@amplitude/mcp-analytics'; import { z, type ZodType } from 'zod'; import { createApiClient, ProfileNotFoundError } from './api-client.js'; import { @@ -14,7 +9,7 @@ import { } from './utils.js'; import { ResponseCache } from './cache.js'; import { AnalyticsHelper } from './analytics.js'; -import { getAmplitudeIdentity } from './amplitude-analytics.js'; +import { setAmplitudeToolContext } from './amplitude-analytics.js'; import type { ApiClient, BrowserlessSession, @@ -165,17 +160,7 @@ export function defineTool( } const apiUrl = s?.apiUrl ?? config.browserlessApiUrl; - if (getCurrentContext()) { - try { - setIdentity({ userId: getAmplitudeIdentity(s, token) }); - if (prompt !== undefined) setRationale(prompt); - } catch (error) { - console.error( - '[browserless-mcp] Amplitude tool context failed:', - error, - ); - } - } + setAmplitudeToolContext(s, token, prompt); def.validateUrl?.(params); diff --git a/test/lib/amplitude-analytics.spec.ts b/test/lib/amplitude-analytics.spec.ts index c383ac1..acd91a3 100644 --- a/test/lib/amplitude-analytics.spec.ts +++ b/test/lib/amplitude-analytics.spec.ts @@ -10,6 +10,8 @@ import { 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 = () => { @@ -126,6 +128,45 @@ describe('Amplitude MCP analytics', () => { 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', @@ -169,4 +210,18 @@ describe('Amplitude MCP analytics', () => { 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); + }); });