From 7c506c027507d3cdb59cca015fca11d16182fce5 Mon Sep 17 00:00:00 2001 From: Evan Wyatt Date: Thu, 22 Oct 2026 13:23:00 -0400 Subject: [PATCH] scope results handles to the caller and label getAccountPlan as JWT --- helius-mcp/src/router/catalog.ts | 9 ++ helius-mcp/src/router/context.ts | 39 +++++-- helius-mcp/src/router/dispatch.ts | 17 +-- helius-mcp/src/utils/request-context.ts | 12 +-- helius-mcp/tests/hosted-eligibility.test.ts | 11 +- helius-mcp/tests/result-ownership.test.ts | 114 ++++++++++++++++++++ 6 files changed, 178 insertions(+), 24 deletions(-) create mode 100644 helius-mcp/tests/result-ownership.test.ts diff --git a/helius-mcp/src/router/catalog.ts b/helius-mcp/src/router/catalog.ts index d90613df..2f2ce996 100644 --- a/helius-mcp/src/router/catalog.ts +++ b/helius-mcp/src/router/catalog.ts @@ -331,6 +331,15 @@ catalog.getAccountStatus = makeEntry('getAccountStatus', { responseFamily: 'record', }); +// Same path as getAccountStatus: the plan, credits, and project come from the +// dashboard session. Without one it can only report "plan unknown", so the +// account-group default of 'none' understated it and let it onto the hosted +// surface. +catalog.getAccountPlan = makeEntry('getAccountPlan', { + ...catalog.getAccountPlan, + authRequirement: 'jwt', +}); + catalog.previewUpgrade = makeEntry('previewUpgrade', { authRequirement: 'jwt', capabilityGate: gate('agent', 'Upgrade preview'), diff --git a/helius-mcp/src/router/context.ts b/helius-mcp/src/router/context.ts index 9418933b..fb528d4b 100644 --- a/helius-mcp/src/router/context.ts +++ b/helius-mcp/src/router/context.ts @@ -1,18 +1,45 @@ -import { randomUUID } from 'node:crypto'; +import { createHash, randomUUID } from 'node:crypto'; import { getNetwork } from '../utils/helius.js'; +import { contextFromExtra, type RequestContext } from '../utils/request-context.js'; -const SESSION_KEY = randomUUID(); +/** + * Owner of results stored for requests that carry no identity: the stdio + * server, where one process serves one user. + */ +const PROCESS_SESSION_KEY = randomUUID(); export type RouterContext = { + /** Owner of the `expandResult` handles this request stores and may redeem. */ sessionKey: string; network: 'mainnet-beta' | 'devnet'; }; -export function getRouterContext(): RouterContext { - const network = getNetwork(); +function transportSessionId(extra: unknown): string | undefined { + if (typeof extra !== 'object' || extra === null) return undefined; + const sessionId = (extra as { sessionId?: unknown }).sessionId; + return typeof sessionId === 'string' && sessionId.length > 0 ? sessionId : undefined; +} + +/** + * Whose handles a request can see. + * + * Without a request context, every request shares the process key, as it + * always has. With one, the owner is the caller's key, narrowed to the + * transport session when the transport has one, so two connections on the + * same API key do not evict or read each other's handles. + */ +function resultOwnerKey(ctx: RequestContext | null, extra: unknown): string { + if (!ctx) return PROCESS_SESSION_KEY; + const sessionId = transportSessionId(extra); + if (!sessionId) return ctx.sessionKey; + return createHash('sha256').update(`${ctx.sessionKey}:${sessionId}`).digest('hex').slice(0, 32); +} +/** Router state for one request. Pass the SDK's `extra` so handles are scoped to the caller. */ +export function getRouterContext(extra?: unknown): RouterContext { + const ctx = contextFromExtra(extra); return { - sessionKey: SESSION_KEY, - network: network === 'devnet' ? 'devnet' : 'mainnet-beta', + sessionKey: resultOwnerKey(ctx, extra), + network: getNetwork(ctx) === 'devnet' ? 'devnet' : 'mainnet-beta', }; } diff --git a/helius-mcp/src/router/dispatch.ts b/helius-mcp/src/router/dispatch.ts index 2274c277..59fbf304 100644 --- a/helius-mcp/src/router/dispatch.ts +++ b/helius-mcp/src/router/dispatch.ts @@ -119,14 +119,14 @@ function buildStoredResult( params: Record, summary: string, text: string, + ownerSessionKey: string, hint?: ContinuationHint, ): StoredResult { - const context = getRouterContext(); const continuation = detectContinuation(params, hint); const sectionHints = collectSectionHints(text); const stored = putStoredResult({ kind: entry.responseFamily, - ownerSessionKey: context.sessionKey, + ownerSessionKey, summary, availableExpansions: buildAvailableExpansions(entry.responseFamily, continuation, sectionHints), payload: { @@ -152,6 +152,7 @@ function normalizeSuccessResponse( text: string, requestedDetail: DetailLevel, allowHandles: boolean, + ownerSessionKey: string, hint?: ContinuationHint, ): RouterResponse { const baseText = text.trim(); @@ -172,7 +173,7 @@ function normalizeSuccessResponse( && (requestedDetail === 'summary' || size > limit); if (needsHandle) { - const stored = buildStoredResult(entry, publicTool, params, summaryText, fullText, hint); + const stored = buildStoredResult(entry, publicTool, params, summaryText, fullText, ownerSessionKey, hint); return mcpResultHandle(summaryText, stored.resultId, stored.availableExpansions, { family: entry.responseFamily, action: entry.action, @@ -194,6 +195,7 @@ function normalizeActionResponse( result: ActionHandlerResponse, requestedDetail: DetailLevel, allowHandles: boolean, + ownerSessionKey: string, ): RouterResponse { const text = toPlainText(result); @@ -205,7 +207,7 @@ function normalizeActionResponse( }); } - return normalizeSuccessResponse(entry, publicTool, params, text, requestedDetail, allowHandles, result.continuation); + return normalizeSuccessResponse(entry, publicTool, params, text, requestedDetail, allowHandles, ownerSessionKey, result.continuation); } async function executeActionViaRouter( @@ -218,8 +220,9 @@ async function executeActionViaRouter( ): Promise { const entry = getActionCatalogEntry(action); try { + const { sessionKey } = getRouterContext(extra); const result = await callActionHandler(action, params, extra); - return normalizeActionResponse(entry, publicTool, params, result, requestedDetail, allowHandles); + return normalizeActionResponse(entry, publicTool, params, result, requestedDetail, allowHandles, sessionKey); } catch (err) { const message = err instanceof Error ? err.message : String(err); return mcpErrorCompact(message, { code: 'HANDLER_ERROR', action, publicTool }); @@ -299,7 +302,8 @@ export async function expandStoredResult( return mcpErrorCompact('Missing required "resultId".', { code: 'MISSING_RESULT_ID' }); } - const context = getRouterContext(); + // Only the caller that stored a handle can redeem it. + const context = getRouterContext(extra); const stored = getStoredResult(resultId, context.sessionKey); if (!stored) { return mcpErrorCompact('Result handle not found or no longer available. Re-run the original action.', { @@ -348,6 +352,7 @@ export async function expandStoredResult( selected, requestedDetail, true, + context.sessionKey, rawResponse.continuation, ); } diff --git a/helius-mcp/src/utils/request-context.ts b/helius-mcp/src/utils/request-context.ts index fca7a0ff..831f5709 100644 --- a/helius-mcp/src/utils/request-context.ts +++ b/helius-mcp/src/utils/request-context.ts @@ -35,16 +35,8 @@ export type RequestContext = { * 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, so `expandResult` handles are bucketed per - * process rather than per caller. - * - * Threading a context through the handlers does not fix that on its own — - * `dispatch.ts` calls `getRouterContext()` for itself, above the handler it - * invokes. Bridging the two belongs with the change that makes the context - * required and removes the globals; until then this field is the declared - * intent and `RouterContext` is the live implementation. + * `getRouterContext(extra)` uses it as the owner of stored results, narrowed + * to the transport session when there is one. */ sessionKey: string; }; diff --git a/helius-mcp/tests/hosted-eligibility.test.ts b/helius-mcp/tests/hosted-eligibility.test.ts index 9de2db71..bb6b098a 100644 --- a/helius-mcp/tests/hosted-eligibility.test.ts +++ b/helius-mcp/tests/hosted-eligibility.test.ts @@ -22,7 +22,7 @@ import { ACTION_NAMES } from '../src/router/actions.js'; describe('the hosted surface', () => { it('is exactly this set', () => { // A snapshot rather than a count: swapping one action for another keeps a - // count at 78 and passes, while this names both the entrant and the leaver + // count at 77 and passes, while this names both the entrant and the leaver // in the failing diff. Update deliberately, and say why in the PR. expect(getHostedActions()).toEqual([ 'accountSubscribe', @@ -31,7 +31,6 @@ describe('the hosted surface', () => { 'deleteWebhook', 'fetchHeliusBlog', 'getAccountInfo', - 'getAccountPlan', 'getAllWebhooks', 'getAsset', 'getAssetProof', @@ -121,6 +120,14 @@ describe('what the boundary excludes, and why', () => { } }); + it('labels getAccountPlan as needing the dashboard session, like getAccountStatus', () => { + // It reads the plan, credits, and project from the session. The + // account-group default of 'none' understated that and made it hosted. + expect(ACTION_CATALOG.getAccountPlan.authRequirement).toBe('jwt'); + expect(ACTION_CATALOG.getAccountPlan.responseFamily).toBe('scalar'); + expect(hostedEligible(ACTION_CATALOG.getAccountPlan)).toBe(false); + }); + it('excludes actions whose credential is the caller\'s but whose effect is not', () => { // These pass the auth check and still have to go: `authRequirement` // describes what a call needs to authenticate, not what it reaches. diff --git a/helius-mcp/tests/result-ownership.test.ts b/helius-mcp/tests/result-ownership.test.ts new file mode 100644 index 00000000..65504cb3 --- /dev/null +++ b/helius-mcp/tests/result-ownership.test.ts @@ -0,0 +1,114 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { dispatchRoutedTool, expandStoredResult } from '../src/router/dispatch.js'; +import { getRouterContext } from '../src/router/context.js'; +import { clearStoredResults } from '../src/results/store.js'; +import { callActionHandler } from '../src/router/action-handlers.js'; + +vi.mock('../src/router/action-handlers.js', async (importOriginal) => { + const original = await importOriginal(); + return { ...original, callActionHandler: vi.fn() }; +}); + +const handler = vi.mocked(callActionHandler); + +/** The SDK `extra` a hosted transport attaches for an authenticated caller. */ +function hostedCaller(apiKey: string, sessionId?: string): Record { + return { authInfo: { token: apiKey, clientId: 'test', scopes: [] }, ...(sessionId ? { sessionId } : {}) }; +} + +const STDIO = {}; +const ADDRESS = '11111111111111111111111111111111'; + +async function storeHandle(extra: unknown): Promise { + handler.mockResolvedValueOnce({ + content: [{ type: 'text', text: 'page 1' }], + continuation: { kind: 'paginationToken', api: 'history', token: 'TKN1' }, + }); + const response = await dispatchRoutedTool( + 'heliusTransaction', + { action: 'getTransactionHistory', address: ADDRESS, detail: 'summary' }, + extra, + ); + const match = response.content[0].text.match(/resultId:\s+([^\n]+)/); + expect(match, 'the summary response should carry a resultId').toBeTruthy(); + return match![1].trim(); +} + +async function expand(resultId: string, extra: unknown, params: Record = {}) { + handler.mockResolvedValueOnce({ + content: [{ type: 'text', text: 'page 2' }], + continuation: { kind: 'paginationToken', api: 'history', token: 'TKN2' }, + }); + return expandStoredResult({ resultId, ...params }, extra); +} + +function wasFound(response: Awaited>): boolean { + return !response.content[0].text.includes('Result handle not found'); +} + +beforeEach(() => { + clearStoredResults(); + handler.mockReset(); +}); + +describe('result handle ownership', () => { + it('works as before for stdio, where requests carry no identity', async () => { + const resultId = await storeHandle(STDIO); + expect(wasFound(await expand(resultId, STDIO))).toBe(true); + }); + + it('keeps one hosted caller from redeeming another caller\'s handle', async () => { + const resultId = await storeHandle(hostedCaller('key-alice')); + + expect(wasFound(await expand(resultId, hostedCaller('key-bob')))).toBe(false); + expect(wasFound(await expand(resultId, hostedCaller('key-alice')))).toBe(true); + }); + + it('does not re-run the stored recipe for a caller who does not own it', async () => { + const resultId = await storeHandle(hostedCaller('key-alice')); + handler.mockClear(); + + await expandStoredResult({ resultId }, hostedCaller('key-bob')); + expect(handler).not.toHaveBeenCalled(); + }); + + it('separates two transport sessions on the same API key', async () => { + const resultId = await storeHandle(hostedCaller('key-alice', 'session-1')); + + expect(wasFound(await expand(resultId, hostedCaller('key-alice', 'session-2')))).toBe(false); + expect(wasFound(await expand(resultId, hostedCaller('key-alice', 'session-1')))).toBe(true); + }); + + it('keeps stdio and hosted callers apart in both directions', async () => { + const stdioHandle = await storeHandle(STDIO); + const hostedHandle = await storeHandle(hostedCaller('key-alice')); + + expect(wasFound(await expand(stdioHandle, hostedCaller('key-alice')))).toBe(false); + expect(wasFound(await expand(hostedHandle, STDIO))).toBe(false); + }); + + it('gives a next-page handle the same owner as the caller who expanded it', async () => { + const first = await storeHandle(hostedCaller('key-alice')); + const next = await expand(first, hostedCaller('key-alice'), { continuation: 'next', detail: 'summary' }); + + const match = next.content[0].text.match(/resultId:\s+([^\n]+)/); + expect(match, 'expanding the next page in summary detail should issue a new handle').toBeTruthy(); + const nextHandle = match![1].trim(); + + expect(wasFound(await expand(nextHandle, hostedCaller('key-bob')))).toBe(false); + expect(wasFound(await expand(nextHandle, hostedCaller('key-alice')))).toBe(true); + }); +}); + +describe('getRouterContext', () => { + it('is stable for the same caller and session', () => { + expect(getRouterContext(hostedCaller('key-alice', 's1')).sessionKey) + .toBe(getRouterContext(hostedCaller('key-alice', 's1')).sessionKey); + expect(getRouterContext().sessionKey).toBe(getRouterContext(STDIO).sessionKey); + }); + + it('does not expose the API key in the owner key', () => { + expect(getRouterContext(hostedCaller('key-alice')).sessionKey).not.toContain('key-alice'); + expect(getRouterContext(hostedCaller('key-alice', 's1')).sessionKey).not.toContain('key-alice'); + }); +});