feat(helius-mcp): Add an Optional Per-Request Context

helius-mcp/src/utils/helius.ts

@@ -4,6 +4,7 @@ import { getSharedApiKey } from './config.js';
 import { wrapClientWithResilience, withResilience, READ_TIMEOUT_MS } from './resilience.js';
 import { registerSecret } from './redact.js';
 import { isSharedCredentialMode } from './runtime.js';
+import type { RequestContext } from './request-context.js';
 
 let sessionApiKey: string | null = null;
 let sessionNetwork: 'mainnet-beta' | 'devnet' = 'mainnet-beta';
@@ -55,8 +56,8 @@ export function setApiKey(apiKey: string): void {
   heliusClient = null; // Reset client so it picks up new key
 }
 
-export function getApiKey(): string {
-  const apiKey = sessionApiKey || process.env.HELIUS_API_KEY || getSharedApiKey();
+export function getApiKey(ctx?: RequestContext): string {
+  const apiKey = ctx?.apiKey || sessionApiKey || process.env.HELIUS_API_KEY || getSharedApiKey();
   if (!apiKey) {
     throw new Error('NO_API_KEY: Set HELIUS_API_KEY environment variable or use setHeliusApiKey tool');
   }
@@ -65,11 +66,20 @@ export function getApiKey(): string {
   return apiKey;
 }
 
-export function hasApiKey(): boolean {
-  return !!(sessionApiKey || process.env.HELIUS_API_KEY || getSharedApiKey());
+export function hasApiKey(ctx?: RequestContext): boolean {
+  return !!(ctx?.apiKey || sessionApiKey || process.env.HELIUS_API_KEY || getSharedApiKey());
 }
 
-export function getHeliusClient(): HeliusClient {
+export function getHeliusClient(ctx?: RequestContext): HeliusClient {
+  // A caller-supplied context gets its own client. Caching per key would be a
+  // cross-tenant hazard for no gain: construction measures ~0.012ms, far below
+  // the network call it wraps.
+  if (ctx) {
+    return wrapClientWithResilience(
+      createHelius({ apiKey: getApiKey(ctx), userAgent: MCP_USER_AGENT }),
+    );
+  }
+
   if (!heliusClient) {
     const apiKey = getApiKey();
     // Wrap so idempotent reads get a timeout + retry-with-backoff; writes,
@@ -84,34 +94,49 @@ export function setNetwork(network: 'mainnet-beta' | 'devnet'): void {
   sessionNetwork = network;
 }
 
-export function getNetwork(): 'mainnet-beta' | 'devnet' {
+export function getNetwork(ctx?: RequestContext): 'mainnet-beta' | 'devnet' {
+  if (ctx) {
+    return ctx.network;
+  }
+
   const envNetwork = process.env.HELIUS_NETWORK;
   if (envNetwork === 'devnet' || envNetwork === 'mainnet-beta') {
     return envNetwork;
   }
   return sessionNetwork;
 }
 
-export function getEnhancedWebSocketUrl(): string {
-  const network = getNetwork();
+export function getEnhancedWebSocketUrl(ctx?: RequestContext): string {
+  const network = getNetwork(ctx);
   const host = network === 'devnet'
     ? 'wss://atlas-devnet.helius-rpc.com'
     : 'wss://atlas-mainnet.helius-rpc.com';
 
-  // Under a shared server-side credential this URL is tool output that reaches
-  // the caller, so it gets a placeholder the caller substitutes with their own
-  // key. Single-tenant stdio still returns a ready-to-use URL.
+  // A supplied context means the key belongs to the caller being answered, so
+  // returning it is the point of this tool, not a leak — check before the
+  // shared-credential placeholder below. Getting this order wrong half-applies
+  // the context: the network comes from the request while the key does not.
+  if (ctx) {
+    return `${host}/?api-key=${getApiKey(ctx)}`;
+  }
+
+  // Without a context the key is the deployment's, not the caller's, so under a
+  // shared credential it gets a placeholder the caller substitutes with their
+  // own. Single-tenant stdio still returns a ready-to-use URL.
   if (isSharedCredentialMode()) {
     return `${host}/?api-key=YOUR_HELIUS_API_KEY`;
   }
 
   return `${host}/?api-key=${getApiKey()}`;
 }
 
-export function getLaserstreamUrl(region?: 'ewr' | 'pitt' | 'slc' | 'lax' | 'lon' | 'ams' | 'fra' | 'tyo' | 'sgp'): string {
+export function getLaserstreamUrl(
+  region?: 'ewr' | 'pitt' | 'slc' | 'lax' | 'lon' | 'ams' | 'fra' | 'tyo' | 'sgp',
+  ctx?: RequestContext,
+): string {
   // Endpoint host is public; clients pass apiKey separately (e.g. @helius/laserstream subscribe options).
   // Do not call getApiKey() here or docs tools like getLaserstreamInfo fail unnecessarily.
-  const network = getNetwork();
+  const network = getNetwork(ctx);
   if (network === 'devnet') {
     return `https://laserstream-devnet-ewr.helius-rpc.com`;
   }
@@ -161,8 +186,12 @@ export async function loadSignerOrFail(): Promise<{ secretKey: Uint8Array; walle
   return { secretKey, walletAddress };
 }
 
-export async function restRequest(endpoint: string, options: RequestInit = {}): Promise<any> {
-  const apiKey = getApiKey();
+export async function restRequest(
+  endpoint: string,
+  options: RequestInit = {},
+  ctx?: RequestContext,
+): Promise<any> {
+  const apiKey = getApiKey(ctx);
   const separator = endpoint.includes('?') ? '&' : '?';
   const url = `https://api.helius.xyz${endpoint}${separator}api-key=${apiKey}`;
 

helius-mcp/src/utils/redact.ts

@@ -16,6 +16,16 @@ const REDACTED = '***REDACTED***';
 /** Matches `api-key=<value>` in a query string, however the URL is delimited. */
 const API_KEY_QUERY_RE = /([?&]api-key=)[^&\s'"`)\]}<]+/gi;
 
+/**
+ * Unbounded and process-global. One stdio caller registers one key, so this is
+ * fine today; a deployment serving many callers accumulates one entry per
+ * distinct key forever, and `redactSecrets` scans the whole set on every
+ * response. There is no cap, no TTL and no eviction.
+ *
+ * Making this request-local belongs with the removal of the other credential
+ * globals, not here — capping it in isolation would silently stop scrubbing an
+ * evicted key, which is worse than the growth.
+ */
 const secrets = new Set<string>();
 
 /**

helius-mcp/src/utils/request-context.ts

@@ -0,0 +1,94 @@
+import { createHash } from 'node:crypto';
+
+/**
+ * Per-request credential and network state.
+ *
+ * Everything the tool layer needs to act *as one caller*. The stdio build has a
+ * single caller for the life of the process, so it resolves this from module
+ * state; a hosted build has a different caller every request and resolves it
+ * from the request itself.
+ *
+ * This type is deliberately small: it carries what identifies and authorises a
+ * caller, not a grab-bag of request metadata. Anything that is the same for
+ * every caller — public docs, the action catalogue — stays shared.
+ *
+ * Lives in `utils/` rather than `router/` because `utils/helius.ts` consumes it
+ * and `utils/` never imports from `router/`.
+ */
+export type RequestContext = {
+  /** The caller's own Helius API key. Downstream calls bill this key. */
+  apiKey: string;
+  /** Helius project the key belongs to, when the transport can resolve it. */
+  projectId?: string;
+  network: 'mainnet-beta' | 'devnet';
+  /**
+   * Opaque, stable per caller — not per request.
+   *
+   * It buckets `expandResult` handles. Per-request would break expansion
+   * immediately, since a handle is issued by one request and redeemed by the
+   * next; per-caller means a handle survives exactly as long as the caller does
+   * and is invisible to everyone else.
+   *
+   * Note this duplicates a responsibility `RouterContext` currently holds:
+   * `getRouterContext()` takes no argument and hands `dispatch.ts` a
+   * process-global session key. Nothing bridges the two yet, and nothing can
+   * until call sites start passing a context. The migration that converts them
+   * is where `getRouterContext` learns to derive from this field; until then
+   * this is the declared intent and `RouterContext` is the live implementation.
+   */
+  sessionKey: string;
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === 'object' && value !== null;
+}
+
+/** Stable, non-reversible bucket id, so the raw key is not also a map key. */
+function deriveSessionKey(apiKey: string, projectId?: string): string {
+  return createHash('sha256')
+    .update(`${projectId ?? ''}:${apiKey}`)
+    .digest('hex')
+    .slice(0, 32);
+}
+
+/**
+ * Build a context from the SDK's per-request `extra`, or `null` when the request
+ * carries no identity.
+ *
+ * `null` is the stdio case and is not an error: callers fall back to the module
+ * state that has always answered for them. A hosted entrypoint is expected to
+ * reject an unidentified request before dispatch ever reaches here.
+ */
+export function contextFromExtra(extra: unknown): RequestContext | null {
+  if (!isRecord(extra)) {
+    return null;
+  }
+
+  const authInfo = extra.authInfo;
+  if (!isRecord(authInfo)) {
+    return null;
+  }
+
+  const apiKey = typeof authInfo.token === 'string' ? authInfo.token : '';
+  if (!apiKey) {
+    return null;
+  }
+
+  // Both of ours live under `extra`, which the SDK documents as the home for
+  // additional token data. Deliberately not `clientId`: that identifies the
+  // OAuth client application, and a Helius project is not one — under any
+  // registration model it would end up holding something else.
+  const authExtra = isRecord(authInfo.extra) ? authInfo.extra : undefined;
+
+  const rawProjectId = authExtra?.projectId;
+  const projectId = typeof rawProjectId === 'string' && rawProjectId ? rawProjectId : undefined;
+
+  const network = authExtra?.network === 'devnet' ? 'devnet' : 'mainnet-beta';
+
+  return {
+    apiKey,
+    projectId,
+    network,
+    sessionKey: deriveSessionKey(apiKey, projectId),
+  };
+}

helius-mcp/tests/request-context.test.ts

@@ -0,0 +1,166 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+
+import { contextFromExtra, type RequestContext } from '../src/utils/request-context.js';
+import { getApiKey, hasApiKey, getNetwork, getEnhancedWebSocketUrl } from '../src/utils/helius.js';
+
+const CALLER_A = 'aaaaaaaa-1111-2222-3333-444444444444';
+const CALLER_B = 'bbbbbbbb-5555-6666-7777-888888888888';
+
+function extraWith(authInfo: unknown): unknown {
+  return { authInfo };
+}
+
+function ctx(overrides: Partial<RequestContext> = {}): RequestContext {
+  return {
+    apiKey: CALLER_A,
+    network: 'mainnet-beta',
+    sessionKey: 'test-session',
+    ...overrides,
+  };
+}
+
+describe('contextFromExtra', () => {
+  it('returns null when the request carries no identity', () => {
+    // stdio: not an error, the caller falls back to module state
+    expect(contextFromExtra(undefined)).toBeNull();
+    expect(contextFromExtra(null)).toBeNull();
+    expect(contextFromExtra('nonsense')).toBeNull();
+    expect(contextFromExtra({})).toBeNull();
+    expect(contextFromExtra(extraWith(undefined))).toBeNull();
+  });
+
+  it('returns null when authInfo carries no token', () => {
+    expect(contextFromExtra(extraWith({ extra: { projectId: 'proj_1' } }))).toBeNull();
+    expect(contextFromExtra(extraWith({ token: '' }))).toBeNull();
+    expect(contextFromExtra(extraWith({ token: 42 }))).toBeNull();
+  });
+
+  it('reads the key from authInfo.token', () => {
+    expect(contextFromExtra(extraWith({ token: CALLER_A }))?.apiKey).toBe(CALLER_A);
+  });
+
+  it('reads the project from authInfo.extra', () => {
+    const built = contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: 'proj_1' } }));
+    expect(built?.projectId).toBe('proj_1');
+  });
+
+  it('leaves the project undefined when it is absent or empty', () => {
+    expect(contextFromExtra(extraWith({ token: CALLER_A }))?.projectId).toBeUndefined();
+    expect(
+      contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: '' } }))?.projectId,
+    ).toBeUndefined();
+  });
+
+  it('ignores clientId, which identifies an OAuth client rather than a project', () => {
+    const built = contextFromExtra(extraWith({ token: CALLER_A, clientId: 'some-oauth-client' }));
+    expect(built?.projectId).toBeUndefined();
+  });
+
+  it('reads the network from authInfo.extra, defaulting to mainnet', () => {
+    expect(contextFromExtra(extraWith({ token: CALLER_A }))?.network).toBe('mainnet-beta');
+    expect(
+      contextFromExtra(extraWith({ token: CALLER_A, extra: { network: 'devnet' } }))?.network,
+    ).toBe('devnet');
+    expect(
+      contextFromExtra(extraWith({ token: CALLER_A, extra: { network: 'nonsense' } }))?.network,
+    ).toBe('mainnet-beta');
+  });
+});
+
+describe('sessionKey', () => {
+  it('is stable across requests from the same caller', () => {
+    // Handles are issued by one request and redeemed by the next, so a
+    // per-request key would break expansion immediately.
+    const first = contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: 'proj_1' } }));
+    const second = contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: 'proj_1' } }));
+    expect(first?.sessionKey).toBe(second?.sessionKey);
+  });
+
+  it('differs between callers', () => {
+    const a = contextFromExtra(extraWith({ token: CALLER_A }));
+    const b = contextFromExtra(extraWith({ token: CALLER_B }));
+    expect(a?.sessionKey).not.toBe(b?.sessionKey);
+  });
+
+  it('differs when the same key is presented under a different project', () => {
+    const a = contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: 'proj_1' } }));
+    const b = contextFromExtra(extraWith({ token: CALLER_A, extra: { projectId: 'proj_2' } }));
+    expect(a?.sessionKey).not.toBe(b?.sessionKey);
+  });
+
+  it('does not contain the raw key', () => {
+    const built = contextFromExtra(extraWith({ token: CALLER_A }));
+    expect(built?.sessionKey).not.toContain(CALLER_A);
+    expect(built?.sessionKey).toHaveLength(32);
+  });
+});
+
+describe('resolvers with a context', () => {
+  const originalKey = process.env.HELIUS_API_KEY;
+  const originalNetwork = process.env.HELIUS_NETWORK;
+
+  beforeEach(() => {
+    process.env.HELIUS_API_KEY = 'env-key-00000000';
+    delete process.env.HELIUS_NETWORK;
+  });
+
+  afterEach(() => {
+    if (originalKey === undefined) delete process.env.HELIUS_API_KEY;
+    else process.env.HELIUS_API_KEY = originalKey;
+    if (originalNetwork === undefined) delete process.env.HELIUS_NETWORK;
+    else process.env.HELIUS_NETWORK = originalNetwork;
+  });
+
+  it('prefers the context key over the environment', () => {
+    expect(getApiKey(ctx())).toBe(CALLER_A);
+  });
+
+  it('falls back to the environment when no context is given', () => {
+    expect(getApiKey()).toBe('env-key-00000000');
+  });
+
+  it('prefers the context network over the environment', () => {
+    process.env.HELIUS_NETWORK = 'mainnet-beta';
+    expect(getNetwork(ctx({ network: 'devnet' }))).toBe('devnet');
+  });
+
+  it('keeps the existing env-then-session order when no context is given', () => {
+    process.env.HELIUS_NETWORK = 'devnet';
+    expect(getNetwork()).toBe('devnet');
+  });
+
+  it('reports a key as present when the context carries one', () => {
+    // Deliberately not asserting the no-context negative: the fallback chain
+    // ends at the on-disk shared config, so that result depends on the machine.
+    delete process.env.HELIUS_API_KEY;
+    expect(hasApiKey(ctx())).toBe(true);
+  });
+
+  it('returns the caller their own key even under a shared credential', () => {
+    // The whole point of this tool is a usable URL. A context means the key
+    // belongs to the caller being answered, so the placeholder must not win.
+    process.env.HELIUS_MCP_SHARED_CREDENTIAL = '1';
+    try {
+      const url = getEnhancedWebSocketUrl(ctx({ network: 'devnet' }));
+      expect(url).toContain(CALLER_A);
+      expect(url).not.toContain('YOUR_HELIUS_API_KEY');
+    } finally {
+      delete process.env.HELIUS_MCP_SHARED_CREDENTIAL;
+    }
+  });
+
+  it('still placeholders the key under a shared credential with no context', () => {
+    process.env.HELIUS_MCP_SHARED_CREDENTIAL = '1';
+    try {
+      expect(getEnhancedWebSocketUrl()).toContain('YOUR_HELIUS_API_KEY');
+    } finally {
+      delete process.env.HELIUS_MCP_SHARED_CREDENTIAL;
+    }
+  });
+
+  it('builds the WebSocket URL from the context, not module state', () => {
+    const url = getEnhancedWebSocketUrl(ctx({ network: 'devnet' }));
+    expect(url).toContain('atlas-devnet');
+    expect(url).toContain(CALLER_A);
+  });
+});