From 34354b548fcf3a58408effe75dda127642de76fb Mon Sep 17 00:00:00 2001
From: Fsocietyhhh <1211904451@qq.com>
Date: Mon, 13 Jul 2026 13:17:08 -0700
Subject: [PATCH] feat(analytics): identify MCP gateway traffic

---
 src/index.ts            |  2 ++
 src/utils/user-agent.ts | 41 ++++++++++++++++++++++++++++++++++++
 test/user-agent.test.ts | 46 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 src/utils/user-agent.ts
 create mode 100644 test/user-agent.test.ts

diff --git a/src/index.ts b/src/index.ts
index 7278b8b..60b505c 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -13,6 +13,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { initializeMcpServer } from "./mcp-handler.js";
 import { warnOnLeakedKeys } from "./utils/key-leak-scanner.js";
+import { installBlockrunMcpUserAgent } from "./utils/user-agent.js";
 import { PROFILES } from "./profiles.js";
 
 // Read version from package.json so it can never drift from the published version.
@@ -54,6 +55,7 @@ function handleCliMetadataFlags(argv: string[]): void {
 }
 
 handleCliMetadataFlags(process.argv);
+installBlockrunMcpUserAgent(VERSION);
 
 async function checkForUpdate() {
   try {
diff --git a/src/utils/user-agent.ts b/src/utils/user-agent.ts
new file mode 100644
index 0000000..61c8bab
--- /dev/null
+++ b/src/utils/user-agent.ts
@@ -0,0 +1,41 @@
+const BLOCKRUN_HOSTS = new Set(["blockrun.ai", "www.blockrun.ai", "sol.blockrun.ai"]);
+const INSTALL_MARKER = Symbol.for("blockrun.mcp.userAgentInstalled");
+
+type MarkedGlobal = typeof globalThis & { [INSTALL_MARKER]?: boolean };
+
+function requestUrl(input: RequestInfo | URL): URL | null {
+  try {
+    const raw = input instanceof Request ? input.url : input.toString();
+    return new URL(raw);
+  } catch {
+    return null;
+  }
+}
+
+/** Wrap fetch so every request to a BlockRun gateway is attributable to MCP. */
+export function withBlockrunMcpUserAgent(
+  baseFetch: typeof fetch,
+  version: string,
+): typeof fetch {
+  const userAgent = `blockrun-mcp/${version}`;
+  return (async (input: RequestInfo | URL, init?: RequestInit) => {
+    const url = requestUrl(input);
+    if (!url || !BLOCKRUN_HOSTS.has(url.hostname)) {
+      return baseFetch(input, init);
+    }
+
+    const headers = new Headers(input instanceof Request ? input.headers : undefined);
+    new Headers(init?.headers).forEach((value, key) => headers.set(key, value));
+    // Deliberately override the SDK's blockrun-ts/* identifier: this process is
+    // the MCP product boundary and should be counted as MCP in gateway logs.
+    headers.set("User-Agent", userAgent);
+    return baseFetch(input, { ...init, headers });
+  }) as typeof fetch;
+}
+
+export function installBlockrunMcpUserAgent(version: string): void {
+  const target = globalThis as MarkedGlobal;
+  if (target[INSTALL_MARKER]) return;
+  globalThis.fetch = withBlockrunMcpUserAgent(globalThis.fetch.bind(globalThis), version);
+  target[INSTALL_MARKER] = true;
+}
diff --git a/test/user-agent.test.ts b/test/user-agent.test.ts
new file mode 100644
index 0000000..0795dc1
--- /dev/null
+++ b/test/user-agent.test.ts
@@ -0,0 +1,46 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { withBlockrunMcpUserAgent } from "../src/utils/user-agent.js";
+
+function recorder() {
+  const calls: Array<{ input: RequestInfo | URL; init?: RequestInit }> = [];
+  const fetcher = (async (input: RequestInfo | URL, init?: RequestInit) => {
+    calls.push({ input, init });
+    return new Response("ok");
+  }) as typeof fetch;
+  return { calls, fetcher };
+}
+
+test("overrides SDK user-agent on BlockRun gateway requests", async () => {
+  const { calls, fetcher } = recorder();
+  const wrapped = withBlockrunMcpUserAgent(fetcher, "0.30.0");
+  await wrapped("https://blockrun.ai/api/v1/chat/completions", {
+    headers: { "User-Agent": "blockrun-ts/3.5.1", "X-Test": "kept" },
+  });
+
+  const headers = new Headers(calls[0].init?.headers);
+  assert.equal(headers.get("user-agent"), "blockrun-mcp/0.30.0");
+  assert.equal(headers.get("x-test"), "kept");
+});
+
+test("marks Solana gateway requests as MCP", async () => {
+  const { calls, fetcher } = recorder();
+  const wrapped = withBlockrunMcpUserAgent(fetcher, "0.30.0");
+  await wrapped("https://sol.blockrun.ai/api/v1/chat/completions");
+  assert.equal(
+    new Headers(calls[0].init?.headers).get("user-agent"),
+    "blockrun-mcp/0.30.0",
+  );
+});
+
+test("does not alter third-party requests", async () => {
+  const { calls, fetcher } = recorder();
+  const wrapped = withBlockrunMcpUserAgent(fetcher, "0.30.0");
+  await wrapped("https://registry.npmjs.org/@blockrun/mcp/latest", {
+    headers: { "User-Agent": "original" },
+  });
+  assert.equal(
+    new Headers(calls[0].init?.headers).get("user-agent"),
+    "original",
+  );
+});