@digital-science-dsl/dimensions-analytics-mcp 1.0.3 → 1.0.4

package/src/client/request-retry.ts

@@ -0,0 +1,115 @@
+/**
+ * Shared retry logic for Dimensions HTTP clients (public API and dsl-service).
+ * @module client/request-retry
+ */
+
+import { NetworkError, RateLimitError, ServerError, TimeoutError } from "./errors.js";
+import type { RateLimiter } from "./rate-limiter.js";
+
+export interface RequestRetryOptions {
+  readonly maxRetries: number;
+  readonly retryDelay: number;
+  readonly rateLimiter?: RateLimiter;
+}
+
+type DelayFn = (ms: number) => Promise<void>;
+
+let delayFn: DelayFn = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+
+/**
+ * Delays before the next retry attempt. Exported for unit tests.
+ */
+export async function requestRetryDelay(ms: number): Promise<void> {
+  return delayFn(ms);
+}
+
+/**
+ * Overrides retry delay behaviour in unit tests.
+ * @internal
+ */
+export function setRequestRetryDelayFn(fn: DelayFn): void {
+  delayFn = fn;
+}
+
+/**
+ * Restores default retry delay behaviour after tests.
+ * @internal
+ */
+export function resetRequestRetryDelayFn(): void {
+  delayFn = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+/**
+ * Returns true when a failed request should be retried.
+ */
+export function isRetryableRequestError(error: unknown): boolean {
+  if (error instanceof RateLimitError) {
+    return true;
+  }
+  if (error instanceof ServerError) {
+    return true;
+  }
+  if (error instanceof NetworkError) {
+    return true;
+  }
+  if (error instanceof TimeoutError) {
+    return true;
+  }
+  return false;
+}
+
+/**
+ * Calculates delay before the next retry (rate-limit aware, exponential backoff with jitter).
+ */
+export function calculateRequestRetryDelay(
+  error: unknown,
+  attempt: number,
+  options: Pick<RequestRetryOptions, "retryDelay" | "rateLimiter">,
+): number {
+  const MAX_DELAY = 60_000;
+
+  if (error instanceof RateLimitError) {
+    const delay =
+      error.clientRateLimit?.retryAfterMs ??
+      options.rateLimiter?.getRetryDelayMs() ??
+      2000 * 2 ** attempt;
+    return Math.min(delay, MAX_DELAY);
+  }
+
+  const backoff = options.retryDelay * 2 ** attempt * (0.5 + Math.random() * 0.5);
+  return Math.min(backoff, MAX_DELAY);
+}
+
+/**
+ * Executes a request with client-side throttling and retries for transient failures.
+ */
+export async function executeRequestWithRetry<T>(
+  requestFn: () => Promise<T>,
+  options: RequestRetryOptions,
+): Promise<T> {
+  let lastError: Error | null = null;
+
+  for (let attempt = 0; attempt <= options.maxRetries; attempt++) {
+    await options.rateLimiter?.waitIfNeeded();
+    try {
+      const result = await requestFn();
+      options.rateLimiter?.recordRequest();
+      return result;
+    } catch (error) {
+      lastError = error as Error;
+
+      if (!isRetryableRequestError(error)) {
+        throw error;
+      }
+
+      if (attempt >= options.maxRetries) {
+        throw error;
+      }
+
+      const retryDelay = calculateRequestRetryDelay(error, attempt, options);
+      await requestRetryDelay(retryDelay);
+    }
+  }
+
+  throw lastError ?? new NetworkError("Unknown error occurred");
+}

package/src/client/usage-context.ts

@@ -0,0 +1,53 @@
+/**
+ * Async-local MCP usage context for dsl-service request headers.
+ * @module client/usage-context
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+
+export type McpDeploymentMode = "hosted" | "local";
+
+export interface McpUsageSession {
+  readonly sessionId: string;
+  readonly deployment: McpDeploymentMode;
+  readonly client: string;
+  readonly version: string;
+}
+
+export interface McpUsageCall {
+  readonly tool: string;
+}
+
+const sessionStorage = new AsyncLocalStorage<McpUsageSession>();
+const callStorage = new AsyncLocalStorage<McpUsageCall>();
+
+let activeSession: McpUsageSession | undefined;
+
+/**
+ * Installs session-level MCP usage metadata for the lifetime of a server instance.
+ */
+export function setMcpUsageSession(session: McpUsageSession): void {
+  activeSession = session;
+}
+
+export function getMcpUsageSession(): McpUsageSession | undefined {
+  return sessionStorage.getStore() ?? activeSession;
+}
+
+/**
+ * Runs a function with session-level usage context (nested over the active session).
+ */
+export function runWithMcpUsageSession<T>(session: McpUsageSession, fn: () => T): T {
+  return sessionStorage.run(session, fn);
+}
+
+/**
+ * Runs an MCP tool handler with per-call usage context (tool name).
+ */
+export function runWithMcpTool<T>(tool: string, fn: () => T): T {
+  return callStorage.run({ tool }, fn);
+}
+
+export function getActiveMcpTool(): string | undefined {
+  return callStorage.getStore()?.tool;
+}

package/src/client/usage-headers.ts

@@ -0,0 +1,45 @@
+/**
+ * X-Dimensions-* usage tracking headers for dsl-service.
+ * @module client/usage-headers
+ */
+
+import { getActiveMcpTool, getMcpUsageSession } from "./usage-context.js";
+
+export const USAGE_HEADER_CHANNEL = "X-Dimensions-Channel";
+export const USAGE_HEADER_MCP_TOOL = "X-Dimensions-MCP-Tool";
+export const USAGE_HEADER_MCP_SESSION_ID = "X-Dimensions-MCP-Session-Id";
+export const USAGE_HEADER_MCP_CLIENT = "X-Dimensions-MCP-Client";
+export const USAGE_HEADER_MCP_VERSION = "X-Dimensions-MCP-Version";
+export const USAGE_HEADER_MCP_DEPLOYMENT = "X-Dimensions-MCP-Deployment";
+
+/**
+ * Builds usage tracking headers from the active MCP session and tool context.
+ */
+export function buildMcpUsageHeaders(): Record<string, string> {
+  const session = getMcpUsageSession();
+  if (!session) {
+    return {};
+  }
+
+  const headers: Record<string, string> = {
+    [USAGE_HEADER_CHANNEL]: "mcp",
+    [USAGE_HEADER_MCP_SESSION_ID]: session.sessionId,
+    [USAGE_HEADER_MCP_CLIENT]: session.client,
+    [USAGE_HEADER_MCP_VERSION]: session.version,
+    [USAGE_HEADER_MCP_DEPLOYMENT]: session.deployment,
+  };
+
+  const tool = getActiveMcpTool();
+  if (tool) {
+    headers[USAGE_HEADER_MCP_TOOL] = tool;
+  }
+
+  return headers;
+}
+
+/**
+ * User-Agent string identifying the MCP server and client app.
+ */
+export function buildMcpUserAgent(version: string, client: string): string {
+  return `dimensions-analytics-mcp/${version} (${client})`;
+}

package/src/dsl/client.ts

@@ -98,15 +98,18 @@ export class DimensionsClient {
 
     if (this.config.backend === "internal") {
       const internal = this.config.internal!;
+      this.rateLimiter = new RateLimiter({
+        maxRequests: 1000,
+        windowMs: 60_000,
+      });
       this.internalClient = new InternalDslClient({
         config: internal.service,
         userEmail: internal.userEmail,
         clientIp: internal.clientIp,
         timeout: this.config.timeout,
-      });
-      this.rateLimiter = new RateLimiter({
-        maxRequests: 1000,
-        windowMs: 60_000,
+        maxRetries: this.config.maxRetries,
+        retryDelay: this.config.retryDelay,
+        rateLimiter: this.rateLimiter,
       });
       return;
     }

package/src/mcp/http-server.ts

@@ -3,6 +3,7 @@
  * @module mcp/http-server
  */
 
+import { randomUUID } from "node:crypto";
 import type { IncomingMessage, Server } from "node:http";
 import { createMcpExpressApp } from "@modelcontextprotocol/sdk/server/express.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
@@ -11,6 +12,7 @@ import type { HostedEnvConfig } from "../client/deployment-config.js";
 import { AuthenticationError } from "../client/errors.js";
 import { parseBearerToken, resolveUserFromApiKey } from "../client/resolve-api-key-user.js";
 import { createMcpServerAsync, warmHostedSchemaCache } from "./server.js";
+import { mcpClientFromUserAgent } from "./usage-tracking.js";
 
 export interface HostedHttpServerOptions {
   readonly hosted: HostedEnvConfig;
@@ -57,11 +59,18 @@ export function startHostedHttpServer(options: HostedHttpServerOptions): HostedH
         clientIp,
       });
 
+      const mcpSessionId = randomUUID();
+      const mcpClient = mcpClientFromUserAgent(
+        typeof req.headers["user-agent"] === "string" ? req.headers["user-agent"] : undefined,
+      );
+
       const { server } = await createMcpServerAsync({
         deploymentMode: "hosted",
         hosted: options.hosted,
         userEmail,
         clientIp,
+        mcpSessionId,
+        mcpClient,
       });
 
       const transport = new StreamableHTTPServerTransport({

package/src/mcp/server.ts

@@ -27,6 +27,7 @@ import { registerLookupTools } from "./tools/lookup.js";
 import { registerQueryTools } from "./tools/query.js";
 import { registerSchemaTools, validateFieldAliases } from "./tools/schema.js";
 import { registerSearchTools } from "./tools/search.js";
+import { initMcpUsageTracking } from "./usage-tracking.js";
 
 /**
  * Configuration options for the MCP server.
@@ -46,6 +47,10 @@ export interface McpServerConfig {
   userEmail?: string;
   /** Client IP to forward to dsl-service and auth.json */
   clientIp?: string;
+  /** MCP session id for usage tracking */
+  mcpSessionId?: string;
+  /** Coarse MCP client label (cursor, vscode, etc.) */
+  mcpClient?: string;
 }
 
 /**
@@ -153,6 +158,7 @@ async function loadSchemaForServer(
  * Warms the shared schema cache for hosted deployment (call once at process startup).
  */
 export async function warmHostedSchemaCache(hosted: HostedEnvConfig): Promise<SchemaStore> {
+  initMcpUsageTracking({ deployment: "hosted", client: "bootstrap" });
   const client = createBootstrapDimensionsClient(hosted);
   return getSharedSchemaStore(client);
 }
@@ -168,6 +174,12 @@ export async function createMcpServerAsync(config: McpServerConfig = {}): Promis
 }> {
   const { mode, hosted } = resolveDeployment(config);
 
+  initMcpUsageTracking({
+    deployment: mode,
+    client: config.mcpClient ?? (mode === "local" ? "stdio" : "unknown"),
+    sessionId: config.mcpSessionId,
+  });
+
   const client = createDimensionsClient({
     mode,
     hosted: hosted ?? config.hosted,

package/src/mcp/tools/analytics.ts

@@ -16,6 +16,7 @@ import {
 } from "../../dsl/index.js";
 import { resolveFacetFieldName, withFieldAliases } from "../middleware/field-aliases.js";
 import type { SchemaStore } from "../schema/index.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import { formatErrorResult, formatToolResult, READ_ONLY_API_ANNOTATIONS } from "../utils.js";
 import { applyAnalyticsFilters } from "./analytics-filters.js";
 
@@ -123,7 +124,8 @@ export function registerAnalyticsTools(
   const entityTypeSchema = entityTypeEnum(schemaStore);
 
   // Facet Query
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "facet_query",
     {
       description:
@@ -212,7 +214,8 @@ export function registerAnalyticsTools(
   );
 
   // Aggregate Query
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "aggregate_query",
     {
       description:
@@ -320,7 +323,8 @@ export function registerAnalyticsTools(
   );
 
   // Citation Trend
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "citation_trend",
     {
       description:
@@ -382,7 +386,8 @@ export function registerAnalyticsTools(
   );
 
   // Funding Trend
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "funding_trend",
     {
       description:

package/src/mcp/tools/fetch-search-pages.ts

@@ -24,6 +24,7 @@ import {
 import { queryHashFromDsl, runAggregateFetch, runFileFetch } from "../batch-fetch.js";
 import { resolveExportFormat } from "../export-format.js";
 import { withFieldAliases } from "../middleware/field-aliases.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import {
   formatErrorResult,
   formatToolResult,
@@ -84,7 +85,8 @@ export function registerFetchSearchPagesTools(
 ): void {
   const entityList = schemaStore.structuredEntityTypes().join(", ");
 
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "fetch_search_pages",
     {
       description:

package/src/mcp/tools/functions.ts

@@ -7,6 +7,7 @@
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z } from "zod";
 import type { DimensionsClient } from "../../dsl/index.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import { formatErrorResult, formatToolResult, READ_ONLY_API_ANNOTATIONS } from "../utils.js";
 
 /**
@@ -15,7 +16,8 @@ import { formatErrorResult, formatToolResult, READ_ONLY_API_ANNOTATIONS } from "
  * @param client - Dimensions client instance
  */
 export function registerFunctionTools(server: McpServer, client: DimensionsClient): void {
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "extract_affiliations",
     {
       description:
@@ -61,7 +63,8 @@ export function registerFunctionTools(server: McpServer, client: DimensionsClien
     },
   );
 
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "extract_grants",
     {
       description:

package/src/mcp/tools/lookup.ts

@@ -8,6 +8,7 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z } from "zod";
 import { type DimensionsClient, EntitySchema, type EntityType } from "../../dsl/index.js";
 import { withFieldAliases } from "../middleware/field-aliases.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import {
   asArray,
   formatErrorResult,
@@ -22,7 +23,8 @@ import {
  */
 export function registerLookupTools(server: McpServer, client: DimensionsClient): void {
   // Get publication by DOI
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "get_by_doi",
     {
       description:
@@ -75,7 +77,8 @@ export function registerLookupTools(server: McpServer, client: DimensionsClient)
   );
 
   // Get publication by PubMed ID
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "get_by_pmid",
     {
       description:
@@ -128,7 +131,8 @@ export function registerLookupTools(server: McpServer, client: DimensionsClient)
   );
 
   // Get entity by Dimensions ID
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "get_by_id",
     {
       description:

package/src/mcp/tools/query.ts

@@ -12,6 +12,7 @@ import {
   validateExecuteDslPolicy,
 } from "../../dsl/index.js";
 import type { SchemaStore } from "../schema/index.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import { formatErrorResult, formatToolResult } from "../utils.js";
 
 /**
@@ -26,7 +27,8 @@ export function registerQueryTools(
   schemaStore: SchemaStore,
 ): void {
   const sources = schemaStore.structuredEntityTypes().join(", ");
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "execute_dsl",
     {
       description: `Execute a raw Dimensions DSL query. Read dimensions://schema/policy and dimensions://examples first. Official docs: https://docs.dimensions.ai/dsl/

package/src/mcp/tools/schema.ts

@@ -9,6 +9,7 @@ import type { DimensionsClient } from "../../dsl/index.js";
 import { buildReverseAliasMap } from "../middleware/field-aliases.js";
 import type { SchemaContext } from "../schema/context.js";
 import { loadSchema, type SchemaStore } from "../schema/index.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import { formatErrorResult, formatToolResult, READ_ONLY_API_ANNOTATIONS } from "../utils.js";
 
 /**
@@ -22,7 +23,8 @@ export function registerSchemaTools(
   client: DimensionsClient,
   context: SchemaContext,
 ): void {
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "describe_schema",
     {
       description:
@@ -98,7 +100,8 @@ export function registerSchemaTools(
     },
   );
 
-  server.registerTool(
+  registerTrackedTool(
+    server,
     "refresh_schema",
     {
       description:

package/src/mcp/tools/search.ts

@@ -19,6 +19,7 @@ import {
   validateSearchPaginationPolicy,
 } from "../../dsl/index.js";
 import { withFieldAliases } from "../middleware/field-aliases.js";
+import { registerTrackedTool } from "../usage-tracking.js";
 import {
   formatErrorResult,
   formatToolResult,
@@ -99,7 +100,8 @@ export function registerSearchTools(
     const resultKey = searchResultKey(source);
     const toolName = searchToolName(source);
 
-    server.registerTool(
+    registerTrackedTool(
+      server,
       toolName,
       {
         description: meta.description,

package/src/mcp/usage-tracking.ts

@@ -0,0 +1,72 @@
+/**
+ * MCP tool usage tracking helpers.
+ * @module mcp/usage-tracking
+ */
+
+import { randomUUID } from "node:crypto";
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import pkg from "../../package.json" with { type: "json" };
+import {
+  type McpDeploymentMode,
+  runWithMcpTool,
+  setMcpUsageSession,
+} from "../client/usage-context.js";
+
+export interface McpUsageTrackingOptions {
+  readonly deployment: McpDeploymentMode;
+  readonly client?: string;
+  readonly sessionId?: string;
+}
+
+/**
+ * Initializes session-level MCP usage metadata for a server instance.
+ */
+export function initMcpUsageTracking(options: McpUsageTrackingOptions): void {
+  setMcpUsageSession({
+    sessionId: options.sessionId ?? randomUUID(),
+    deployment: options.deployment,
+    client: options.client ?? "unknown",
+    version: pkg.version,
+  });
+}
+
+/**
+ * Wraps a tool handler so the active tool name is available for usage headers.
+ */
+export function withMcpUsageTracking<TArgs, TResult>(
+  toolName: string,
+  handler: (args: TArgs) => TResult | Promise<TResult>,
+): (args: TArgs) => Promise<TResult> {
+  return async (args: TArgs) => runWithMcpTool(toolName, () => handler(args));
+}
+
+/**
+ * Registers an MCP tool with usage tracking for dsl-service headers.
+ */
+export function registerTrackedTool(
+  server: McpServer,
+  toolName: string,
+  config: Record<string, unknown>,
+  // biome-ignore lint/suspicious/noExplicitAny: preserve MCP SDK tool handler inference
+  handler: (args: any) => any,
+): void {
+  server.registerTool(
+    toolName,
+    config,
+    withMcpUsageTracking(toolName, handler) as Parameters<McpServer["registerTool"]>[2],
+  );
+}
+
+/**
+ * Derives a coarse MCP client label from an HTTP User-Agent header.
+ */
+export function mcpClientFromUserAgent(userAgent: string | undefined): string {
+  if (!userAgent) return "unknown";
+  const ua = userAgent.toLowerCase();
+  if (ua.includes("cursor")) return "cursor";
+  if (ua.includes("claude")) return "claude-desktop";
+  if (ua.includes("vscode") || ua.includes("visual studio code")) return "vscode";
+  if (ua.includes("windsurf")) return "windsurf";
+  if (ua.includes("copilot")) return "copilot";
+  return "unknown";
+}

package/test/client/deployment-config.test.ts

@@ -4,12 +4,7 @@
  */
 
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
-import {
-  buildHostedDslLoggingInfo,
-  hostedBootstrapUser,
-  loadDeploymentConfig,
-  mcpTrackingUser,
-} from "../../src/client/deployment-config.js";
+import { hostedBootstrapUser, loadDeploymentConfig } from "../../src/client/deployment-config.js";
 
 const HOSTED_ENV = {
   DEPLOYMENT_MODE: "hosted",
@@ -72,22 +67,3 @@ describe("hostedBootstrapUser", () => {
     expect(hostedBootstrapUser("app.dimensions.ai")).toBe("mcp@app.dimensions.ai");
   });
 });
-
-describe("mcpTrackingUser", () => {
-  it("prefixes canonical email with mcp+", () => {
-    expect(mcpTrackingUser("alice@university.edu")).toBe("mcp+alice@university.edu");
-  });
-});
-
-describe("buildHostedDslLoggingInfo", () => {
-  it("includes channel metadata for dsl-service logs", () => {
-    expect(buildHostedDslLoggingInfo("alice@university.edu", "standard")).toEqual({
-      user: "alice@university.edu",
-      dimensions_user: "alice@university.edu",
-      channel: "mcp",
-      mcp_user: "mcp+alice@university.edu",
-      product_variant: "standard",
-      source: "dimensions-analytics-mcp-hosted",
-    });
-  });
-});