From d0fff63af7c84436a9ae131c267e679351252810 Mon Sep 17 00:00:00 2001 From: Rohit Gupta Date: Tue, 23 Jun 2026 13:35:47 +0530 Subject: [PATCH 1/2] SCAL-319970 Add OAuth-gated list_orgs tool Add a new `list_orgs` MCP tool that lists the ACTIVE Orgs configured on the ThoughtSpot instance, available only to connections authenticated via OAuth. Auth method is per client connection (decided by the endpoint the client connects to), not per user or per deployment. To gate on it at runtime, this introduces an explicit `authMode` field on the connection `Props`: - OAuth routes (/mcp, /sse) set authMode = "oauth" - Static-token routes (/bearer/*, /token/*) set authMode = "bearer" | "token" The tool is gated with defense in depth: - Filtered out of listTools() for non-OAuth connections - Rejected in callTool() if invoked directly by a non-OAuth connection Data comes from the public POST /api/rest/2.0/orgs/search endpoint (client.searchOrgs), filtered to ACTIVE orgs. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/bearer.ts | 3 ++ src/index.ts | 2 ++ src/metrics/runtime/tool-metrics.ts | 1 + src/servers/mcp-server.ts | 41 +++++++++++++++++++++++++- src/servers/tool-definitions.ts | 32 ++++++++++++++++++++ src/thoughtspot/thoughtspot-service.ts | 34 +++++++++++++++++++++ src/thoughtspot/types.ts | 7 +++++ src/utils.ts | 8 ++++- test/bearer.spec.ts | 2 ++ 9 files changed, 128 insertions(+), 2 deletions(-) diff --git a/src/bearer.ts b/src/bearer.ts index 9b9478f..51f6846 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -91,6 +91,9 @@ async function handleTokenAuth( accessToken: accessToken, instanceUrl: validateAndSanitizeUrl(tsHost), clientName, + // Distinguishes static-token auth ("bearer"/"token") from OAuth. Used to gate + // OAuth-only tools such as `list_orgs`. + authMode: authRouteFamily, }; const requestedApiVersion = url.searchParams.get("api-version"); diff --git a/src/index.ts b/src/index.ts index dbec08b..35a98c5 100644 --- a/src/index.ts +++ b/src/index.ts @@ -72,6 +72,8 @@ function createMCPRouter( ? normalizeRequestedApiVersionForAnalytics(requestedApiVersion) : undefined, apiVersionMode, + // These routers are only mounted for the OAuth endpoints (/mcp, /sse). + authMode: "oauth", }; // Route to the appropriate serve method diff --git a/src/metrics/runtime/tool-metrics.ts b/src/metrics/runtime/tool-metrics.ts index c86b7ee..5ca4070 100644 --- a/src/metrics/runtime/tool-metrics.ts +++ b/src/metrics/runtime/tool-metrics.ts @@ -21,6 +21,7 @@ export const UPSTREAM_OPERATION_NAMES = { "send_agent_conversation_message_streaming", importMetadataTml: "import_metadata_tml", searchMetadata: "search_metadata", + searchOrgs: "search_orgs", } as const; export type UpstreamOperation = diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 334b551..0c18175 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -25,6 +25,7 @@ import { GetDataSourceSuggestionsSchema, GetRelevantQuestionsSchema, GetSessionUpdatesInputSchema, + ListOrgsInputSchema, SendSessionMessageInputSchema, ToolName, } from "./tool-definitions"; @@ -43,6 +44,14 @@ export class MCPServer extends BaseMCPServer { return "mcp"; } + /** + * Whether the current connection authenticated via OAuth (as opposed to a static + * bearer/token). Used to gate OAuth-only tools such as `list_orgs`. + */ + protected isOAuthAuth(): boolean { + return this.ctx.props.authMode === "oauth"; + } + protected getToolMetricApiVersionLabel(): string | undefined { const apiVersion = this.ctx.props.apiVersion; if (typeof apiVersion !== "string" || apiVersion.length === 0) { @@ -136,6 +145,11 @@ export class MCPServer extends BaseMCPServer { ); } + // Org tools (e.g. list_orgs) are only available to OAuth-authenticated users. + if (!this.isOAuthAuth()) { + tools = tools.filter((tool) => tool.name !== ToolName.ListOrgs); + } + return { tools }; } @@ -192,7 +206,7 @@ export class MCPServer extends BaseMCPServer { switch (name) { case ToolName.Ping: { if (this.ctx.props.accessToken && this.ctx.props.instanceUrl) { - if (!this.getThoughtSpotService(recorder).validateConnection()) { + if (!this.getThoughtSpotService(recorder).validateConnection()) { return this.createErrorResponse( "Failed to validate connection", "Ping failed", @@ -254,6 +268,18 @@ export class MCPServer extends BaseMCPServer { return this.callCreateDashboard(request, recorder); } + case ToolName.ListOrgs: { + // Defense in depth: this tool is omitted from listTools for non-OAuth + // connections, but reject direct invocation as well. + if (!this.isOAuthAuth()) { + return this.createErrorResponse( + "The list_orgs tool is only available when authenticated via OAuth.", + "List orgs rejected: non-OAuth auth mode", + ); + } + return this.callListOrgs(recorder); + } + default: throw new Error(`Unknown tool: ${name}`); } @@ -612,6 +638,19 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; ); } + @WithSpan("call-list-orgs") + async callListOrgs(recorder: MetricsRecorder) { + const span = trace.getSpan(context.active()); + + const orgs = await this.getThoughtSpotService(recorder).searchOrgs(); + span?.setAttribute("total_orgs", orgs.length); + + return this.createStructuredContentSuccessResponse( + { orgs }, + `${orgs.length} org(s) found`, + ); + } + private _sources: { list: DataSource[]; map: Map; diff --git a/src/servers/tool-definitions.ts b/src/servers/tool-definitions.ts index 8db2e91..f2c6f96 100644 --- a/src/servers/tool-definitions.ts +++ b/src/servers/tool-definitions.ts @@ -257,6 +257,24 @@ export const CreateDashboardOutputSchema = z.object({ ), }); +export const ListOrgsInputSchema = z.object({}); + +export const ListOrgsOutputSchema = z.object({ + orgs: z + .array( + z.object({ + id: z.number().describe("Unique identifier of the Org."), + name: z.string().describe("Name of the Org."), + description: z.string().optional().describe("Description of the Org."), + status: z + .string() + .optional() + .describe("Status of the Org (ACTIVE or IN_ACTIVE)."), + }), + ) + .describe("The list of Orgs the user can access."), +}); + export enum ToolName { // V1 Ping = "ping", @@ -270,6 +288,7 @@ export enum ToolName { SendSessionMessage = "send_session_message", GetSessionUpdates = "get_session_updates", CreateDashboard = "create_dashboard", + ListOrgs = "list_orgs", } export const toolDefinitionsV1 = [ @@ -398,4 +417,17 @@ export const toolDefinitionsV2 = [ openWorldHint: false, }, }, + { + name: ToolName.ListOrgs, + description: + "List the Orgs configured on the ThoughtSpot instance, including the ID, name, description, and status of each Org, along with the ID of the Org that is currently active for this session. Only available when authenticated via OAuth.", + inputSchema: z.toJSONSchema(ListOrgsInputSchema), + outputSchema: z.toJSONSchema(ListOrgsOutputSchema), + annotations: { + title: "List Orgs", + readOnlyHint: true, + destructiveHint: false, + openWorldHint: false, + }, + }, ]; diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 974cb17..2a4b19a 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -26,6 +26,7 @@ import type { DataSource, DataSourceSuggestion, Message, + Org, SessionInfo, } from "./types"; @@ -717,6 +718,39 @@ export class ThoughtSpotService { }; } + /** + * List the Orgs configured on the ThoughtSpot instance. + * + * Uses the public POST /api/rest/2.0/orgs/search endpoint. Requires Org + * administration privileges; callers without them will receive an upstream + * 401/403, which is surfaced to the tool layer. + */ + @WithSpan("search-orgs") + async searchOrgs(): Promise { + const span = getActiveSpan(); + + const orgs = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.searchOrgs, + () => this.client.searchOrgs({ status: "ACTIVE" }), + ); + + const results: Org[] = (orgs ?? []) + .filter( + (org): org is typeof org & { id: number } => + typeof org.id === "number" && org.status === "ACTIVE", + ) + .map((org) => ({ + id: org.id, + name: org.name ?? "", + description: org.description ?? undefined, + status: org.status ?? undefined, + })); + + span?.setAttribute("results_count", results.length); + + return results; + } + /** * Search worksheets by term */ diff --git a/src/thoughtspot/types.ts b/src/thoughtspot/types.ts index b4d6a60..cca4f1d 100644 --- a/src/thoughtspot/types.ts +++ b/src/thoughtspot/types.ts @@ -18,6 +18,13 @@ export interface DataSourceSuggestionResponse { dataSources: DataSourceSuggestion[]; } +export interface Org { + id: number; + name: string; + description?: string; + status?: string; +} + export interface SessionInfo { mixpanelToken: string; userGUID: string; diff --git a/src/utils.ts b/src/utils.ts index b7c525d..0dd85d5 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1,5 +1,5 @@ import { type Span, SpanStatusCode } from "@opentelemetry/api"; -import type { ApiVersionMode } from "./metrics/runtime/metric-types"; +import type { ApiVersionMode, AuthMode } from "./metrics/runtime/metric-types"; import { getActiveSpan } from "./metrics/tracing/tracing-utils"; export type Props = { @@ -13,6 +13,12 @@ export type Props = { apiVersion?: string; apiVersionMode?: ApiVersionMode; apiRequestedVersion?: string; + /** + * How the user authenticated for this connection. Used to gate tools that are only + * available to OAuth users (e.g. `list_orgs`). Set to "oauth" by the OAuth flow and + * "bearer"/"token" by the static-token flow in bearer.ts. + */ + authMode?: AuthMode; }; export class McpServerError extends Error { diff --git a/test/bearer.spec.ts b/test/bearer.spec.ts index a0af829..debb7f9 100644 --- a/test/bearer.spec.ts +++ b/test/bearer.spec.ts @@ -184,6 +184,7 @@ describe("Bearer Handler", () => { clientName: "Custom Test Client", apiVersion: "backwards-compatibility-default", apiVersionMode: "implicit_legacy", + authMode: "bearer", }); // Verify the response @@ -210,6 +211,7 @@ describe("Bearer Handler", () => { clientName: "Bearer Token client", apiVersion: "backwards-compatibility-default", apiVersionMode: "implicit_legacy", + authMode: "bearer", }); // Verify the response From ae25ec47ae7589890fe0ebf9ad371be2e032afab Mon Sep 17 00:00:00 2001 From: Rohit Gupta Date: Wed, 24 Jun 2026 23:40:25 +0530 Subject: [PATCH 2/2] SCAL-319970 Add switch_org + durable org context, gate org tools on Orgs-enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Build out multi-org support on top of the OAuth-gated list_orgs tool. Tools - Add switch_org: mints an org-scoped bearer token for the requested org and makes subsequent tool calls operate against it. No pre-validation against list_orgs — a 401 from minting is surfaced as "org not accessible". - Gate both list_orgs and switch_org on OAuth AND cluster Orgs-enabled (configInfo.orgsConfiguration.enabled from session info), failing closed. Hidden from listTools and rejected in callTool when unavailable. Token flow - OAuth callback fetches the global token via session/v2/gettoken?refresh=true and stores token + refreshToken + tokenCreatedTime + tokenExpiryDuration. - getRefreshedToken sends the refresh token via the X-Refresh-Token header (verified this is what mints a fresh access token). - Org-scoped tokens are minted on demand from the global token and cached in-memory per instance. Durable active org (multi-session safe) - The active org is stored in a per-user Durable Object instance keyed by hash(refreshToken ?? accessToken) via the existing conversation-storage DO namespace. The refresh-token hash is stable across the access token's 24h rotation and shared across the multiple MCP sessions/DOs a single client opens, so a switch in one session is visible to the others. Conversation storage is re-keyed onto the same hash so it survives token refresh too. - The active org is NOT stored in props (props are rebuilt from the OAuth grant on every request and would clobber it). Tests - Org tool gate (OAuth + orgs-enabled, fail-closed), getStorageKeyHash (refresh-token preference + access-token fallback), isOrgsEnabled, searchOrgs ACTIVE filtering, fetchOrgBearerToken/getRefreshedToken delegation, switch_org persistence + 401 path, and shared active-org store across instances. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/handlers.ts | 3 + src/metrics/runtime/tool-metrics.ts | 2 + src/servers/conversation-storage-server.ts | 21 ++ src/servers/mcp-server-base.ts | 102 +++++++- src/servers/mcp-server.ts | 204 ++++++++++++++- src/servers/tool-definitions.ts | 36 ++- src/storage-service/storage-service.ts | 35 +++ src/thoughtspot/thoughtspot-client.ts | 206 ++++++++++++--- src/thoughtspot/thoughtspot-service.ts | 34 +++ src/thoughtspot/types.ts | 3 + src/utils.ts | 16 ++ static/oauth-callback.js | 17 +- test/servers/mcp-server-base.spec.ts | 96 ++++++- test/servers/mcp-server-orgs.spec.ts | 261 +++++++++++++++++++ test/thoughtspot/thoughtspot-service.spec.ts | 53 ++++ 15 files changed, 1025 insertions(+), 64 deletions(-) create mode 100644 test/servers/mcp-server-orgs.spec.ts diff --git a/src/handlers.ts b/src/handlers.ts index 994fe6e..d557264 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -260,6 +260,9 @@ class Handler { scope: oauthReqInfo.scope, props: { accessToken: token.data.token, + refreshToken: token.data.refreshToken, + tokenCreatedTime: token.data.tokenCreatedTime, + tokenExpiryDuration: token.data.tokenExpiryDuration, instanceUrl: instanceUrl, clientName: clientName, } as Props, diff --git a/src/metrics/runtime/tool-metrics.ts b/src/metrics/runtime/tool-metrics.ts index 5ca4070..f0a637c 100644 --- a/src/metrics/runtime/tool-metrics.ts +++ b/src/metrics/runtime/tool-metrics.ts @@ -22,6 +22,8 @@ export const UPSTREAM_OPERATION_NAMES = { importMetadataTml: "import_metadata_tml", searchMetadata: "search_metadata", searchOrgs: "search_orgs", + fetchOrgBearerToken: "fetch_org_bearer_token", + getRefreshedToken: "get_refreshed_token", } as const; export type UpstreamOperation = diff --git a/src/servers/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts index a4f8b4f..b4a32e6 100644 --- a/src/servers/conversation-storage-server.ts +++ b/src/servers/conversation-storage-server.ts @@ -6,6 +6,8 @@ const STORAGE_BATCH_SIZE = 127; // Cloudflare DO bulk get/put limit is 128, we u const MESSAGE_KEY_PREFIX = "message-"; const IS_DONE_KEY = "is-done"; +// Storage key for the active org on the per-user active-org DO instance. +const ACTIVE_ORG_KEY = "active-org"; const WRITE_BOOKMARK_KEY = "write-bookmark"; const READ_BOOKMARK_KEY = "read-bookmark"; @@ -56,6 +58,25 @@ export class ConversationStorageServerSQLite { return Response.json(state); } + // Active-org state. Stored on a DO addressed by the user's storage-key + // hash (not a conversation id), so it is shared across all of the + // user's MCP sessions. No TTL: it must persist until the user switches + // again or reauthenticates (a new login yields a new hash). + case "GET /active-org": { + const activeOrgId = + (await this.state.storage.get(ACTIVE_ORG_KEY)) ?? null; + return Response.json({ activeOrgId }); + } + + case "POST /active-org": { + const body = (await request.json()) as { activeOrgId: string }; + await this.state.storage.put( + ACTIVE_ORG_KEY, + body.activeOrgId, + ); + return Response.json({ ok: true }); + } + default: return new Response("Not Found", { status: 404 }); } diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 4df14f0..15221c9 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -93,6 +93,18 @@ export abstract class BaseMCPServer extends Server { return String(this.sessionInfo.enableSpotterDataSourceDiscovery) === "true"; } + /** + * Whether Orgs are enabled on this cluster (from session info). Fails closed: + * if session info is unavailable or the flag is absent, returns false so the + * org tools stay hidden. + */ + protected isOrgsEnabled(): boolean { + if (!this.sessionInfo) { + return false; + } + return this.sessionInfo.orgsEnabled === true; + } + /** * Initialize span with common attributes (user_guid and instance_url) */ @@ -187,19 +199,30 @@ export abstract class BaseMCPServer extends Server { }; } - protected async getStorageService(): Promise { - const accessToken = this.ctx.props.accessToken; - if (!accessToken || accessToken.length === 0) { - throw new Error("Access token is required to use Storage Service"); + /** + * Stable per-login hash used to namespace this user's durable storage (both + * conversation buffers and active-org state), keeping users isolated. + * + * Keyed on the refresh token when present (OAuth): it is stable across the + * access token's 24h rotation and only changes on full reauthentication, so + * storage survives token refresh and resets on reauth. Falls back to the + * access token for static bearer/token connections, which have no refresh + * token (their token is long-lived). + */ + protected async getStorageKeyHash(): Promise { + const keyToken = this.ctx.props.refreshToken ?? this.ctx.props.accessToken; + if (!keyToken || keyToken.length === 0) { + throw new Error("A token is required to derive the storage key"); } - const encodedAccessToken = new TextEncoder().encode(accessToken); const hashBuffer = await crypto.subtle.digest( "SHA-256", - encodedAccessToken, - ); - const hashUrlSafe = Buffer.from(new Uint8Array(hashBuffer)).toString( - "base64url", + new TextEncoder().encode(keyToken), ); + return Buffer.from(new Uint8Array(hashBuffer)).toString("base64url"); + } + + protected async getStorageService(): Promise { + const hashUrlSafe = await this.getStorageKeyHash(); return new StorageServiceClient( this.ctx.env .CONVERSATION_STORAGE_OBJECT as unknown as DurableObjectNamespace, @@ -207,6 +230,50 @@ export abstract class BaseMCPServer extends Server { ); } + /** + * The org currently active for this session, if any. When set, all + * ThoughtSpot calls are scoped to this org via the x-thoughtspot-orgs header. + * Subclasses override this to expose their per-session org state. Defaults to + * undefined (the user's default org, as resolved by the cluster). + */ + protected getActiveOrgId(): string | undefined { + return undefined; + } + + /** + * The bearer token to use for ThoughtSpot calls. Defaults to the token from + * the session. Subclasses override this to return an org-scoped bearer token + * when an org has been selected. + */ + protected getActiveBearerToken(): string { + return this.ctx.props.accessToken; + } + + /** + * Build a ThoughtSpot service bound to an explicit bearer token and org, + * bypassing the active-org/token resolution. Used for org-token minting, + * which must authenticate with a specific token. + */ + protected getThoughtSpotServiceWithToken( + bearerToken: string, + orgId?: string, + recorder?: MetricsRecorder, + analyticsContextOverride?: MetricAnalyticsContext, + ) { + return new ThoughtSpotService( + getThoughtSpotClient(this.ctx.props.instanceUrl, bearerToken, orgId), + { + recorder, + metricsEnv: this.ctx.env as unknown as Record, + waitUntil: this.getMetricsWaitUntil(), + analyticsContext: this.mergeMetricAnalyticsContext( + analyticsContextOverride, + ), + eventIdentity: this.getMetricEventIdentity(), + }, + ); + } + protected getThoughtSpotService( recorder?: MetricsRecorder, analyticsContextOverride?: MetricAnalyticsContext, @@ -214,7 +281,8 @@ export abstract class BaseMCPServer extends Server { return new ThoughtSpotService( getThoughtSpotClient( this.ctx.props.instanceUrl, - this.ctx.props.accessToken, + this.getActiveBearerToken(), + this.getActiveOrgId(), ), { recorder, @@ -439,8 +507,22 @@ export abstract class BaseMCPServer extends Server { }); }, ); + + // Subclass post-initialization hook (runs after sessionInfo is available + // and handlers are registered). Best-effort: failures must not break the + // connection. + try { + await this.postInit(); + } catch (error) { + console.error("postInit failed:", error); + } } + /** + * Optional hook for subclasses to run setup after init(). Default no-op. + */ + protected async postInit(): Promise {} + async addTracker(tracker: Tracker) { this.trackers.add(tracker); } diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 0c18175..859b1ab 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -25,8 +25,8 @@ import { GetDataSourceSuggestionsSchema, GetRelevantQuestionsSchema, GetSessionUpdatesInputSchema, - ListOrgsInputSchema, SendSessionMessageInputSchema, + SwitchOrgInputSchema, ToolName, } from "./tool-definitions"; import { @@ -36,10 +36,116 @@ import { } from "./version-registry"; export class MCPServer extends BaseMCPServer { + // In-memory mirror of the active org, loaded once per request lifecycle from + // the shared per-user store (keyed by the storage-key hash, so it is the same + // across all of the user's MCP sessions/DOs). Read synchronously by the + // active-org/token accessors; the durable source of truth is the store. + private activeOrgId: string | undefined; + + // Org-scoped bearer tokens, keyed by org id. In-memory only (runtime-derived, + // not identity); re-minted as needed. + private orgBearerTokens: Record = {}; + constructor(ctx: Context) { super(ctx, "ThoughtSpot", "2.0.0"); } + /** + * The single accessor for the active org. Reads the in-memory mirror, which is + * loaded from the shared store on connect (postInit). + */ + protected getActiveOrgId(): string | undefined { + return this.activeOrgId; + } + + /** + * Use the org-scoped bearer token for the active org if we hold one; otherwise + * fall back to the session's global access token (from props/grant). + */ + protected getActiveBearerToken(): string { + const orgToken = this.activeOrgId + ? this.orgBearerTokens[this.activeOrgId] + : undefined; + return orgToken ?? this.ctx.props.accessToken; + } + + /** + * Load the active org from the shared per-user store into the in-memory mirror. + * Keyed by the storage-key hash (refresh-token based), so a switch made in any + * of the user's MCP sessions is visible here. + */ + private async loadActiveOrg(): Promise { + const storage = await this.getStorageService(); + const stored = await storage.getActiveOrg(); + // Always reflect the store (including back to undefined if cleared), since + // the value may have changed in another of the user's sessions/DOs. + this.activeOrgId = stored ?? undefined; + } + + /** + * Re-read the active org from the shared store. Because the MCP client may + * fan requests across multiple DOs, we re-read on each org-aware tool call so + * a switch made elsewhere is reflected, rather than trusting a stale mirror. + */ + private async ensureActiveOrgLoaded(): Promise { + await this.loadActiveOrg(); + } + + /** + * Persist the active org to the shared per-user store and update the in-memory + * mirror. Shared across the user's sessions; persists until the next switch or + * reauthentication (a new login changes the storage-key hash). + */ + private async setActiveOrg(orgId: string): Promise { + this.activeOrgId = orgId; + const storage = await this.getStorageService(); + await storage.setActiveOrg(orgId); + } + + /** + * Ensure we hold an org-scoped bearer token for `orgId`: reuse the in-memory + * one, or mint a new one from the global access token. Returns the org token. + */ + private async ensureOrgToken( + orgId: string, + recorder?: MetricsRecorder, + ): Promise { + const cached = this.orgBearerTokens[orgId]; + if (cached) { + return cached; + } + const orgToken = await this.getThoughtSpotServiceWithToken( + this.ctx.props.accessToken, + undefined, + recorder, + ).fetchOrgBearerToken(this.ctx.props.accessToken, orgId); + this.orgBearerTokens[orgId] = orgToken; + return orgToken; + } + + /** + * On connect: load the active org from the shared per-user store. If one is + * stored (the user switched in this or another session), mint its token so + * subsequent calls are scoped to it. If none is stored, do NOT write one — the + * absence means "no explicit switch", and calls fall back to the global token / + * the cluster-resolved default org. Best-effort: failures must not break the + * connection. Only OAuth sessions can mint org tokens, so this is a no-op + * otherwise. + */ + protected async postInit(): Promise { + if (!this.isOAuthAuth()) { + return; + } + try { + await this.loadActiveOrg(); + if (this.activeOrgId) { + await this.ensureOrgToken(this.activeOrgId); + } + } catch (error) { + console.error("Failed to load/mint active org on connect:", error); + } + } + protected getToolMetricApiSurface(): ToolMetricApiSurface { return "mcp"; } @@ -52,6 +158,15 @@ export class MCPServer extends BaseMCPServer { return this.ctx.props.authMode === "oauth"; } + /** + * Org tools (list_orgs/switch_org) are available only when the connection is + * OAuth (the only auth mode that can mint org-scoped tokens) AND the cluster + * has Orgs enabled. Fails closed if either is unknown. + */ + protected areOrgToolsAvailable(): boolean { + return this.isOAuthAuth() && this.isOrgsEnabled(); + } + protected getToolMetricApiVersionLabel(): string | undefined { const apiVersion = this.ctx.props.apiVersion; if (typeof apiVersion !== "string" || apiVersion.length === 0) { @@ -145,9 +260,13 @@ export class MCPServer extends BaseMCPServer { ); } - // Org tools (e.g. list_orgs) are only available to OAuth-authenticated users. - if (!this.isOAuthAuth()) { - tools = tools.filter((tool) => tool.name !== ToolName.ListOrgs); + // Org tools (list_orgs, switch_org) require OAuth AND Orgs enabled on the + // cluster. + if (!this.areOrgToolsAvailable()) { + tools = tools.filter( + (tool) => + tool.name !== ToolName.ListOrgs && tool.name !== ToolName.SwitchOrg, + ); } return { tools }; @@ -269,17 +388,29 @@ export class MCPServer extends BaseMCPServer { } case ToolName.ListOrgs: { - // Defense in depth: this tool is omitted from listTools for non-OAuth - // connections, but reject direct invocation as well. - if (!this.isOAuthAuth()) { + // Defense in depth: omitted from listTools when unavailable, but + // reject direct invocation as well. + if (!this.areOrgToolsAvailable()) { return this.createErrorResponse( - "The list_orgs tool is only available when authenticated via OAuth.", - "List orgs rejected: non-OAuth auth mode", + "The list_orgs tool is only available when authenticated via OAuth on a cluster with Orgs enabled.", + "List orgs rejected: org tools unavailable", ); } return this.callListOrgs(recorder); } + case ToolName.SwitchOrg: { + // Defense in depth: omitted from listTools when unavailable, but + // reject direct invocation as well. + if (!this.areOrgToolsAvailable()) { + return this.createErrorResponse( + "The switch_org tool is only available when authenticated via OAuth on a cluster with Orgs enabled.", + "Switch org rejected: org tools unavailable", + ); + } + return this.callSwitchOrg(request, recorder); + } + default: throw new Error(`Unknown tool: ${name}`); } @@ -645,12 +776,65 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; const orgs = await this.getThoughtSpotService(recorder).searchOrgs(); span?.setAttribute("total_orgs", orgs.length); + // Read the active org from the shared per-user store (the single source of + // truth). Ensure it's loaded first — this request may run on a different DO + // than the one that handled postInit/switch_org. + await this.ensureActiveOrgLoaded(); + const activeOrgId = this.getActiveOrgId(); + return this.createStructuredContentSuccessResponse( - { orgs }, + { + orgs: orgs.map((org) => ({ + ...org, + is_active: + activeOrgId !== undefined && String(org.id) === activeOrgId, + })), + }, `${orgs.length} org(s) found`, ); } + @WithSpan("call-switch-org") + async callSwitchOrg( + request: z.infer, + recorder: MetricsRecorder, + ) { + const span = trace.getSpan(context.active()); + const { org_id } = SwitchOrgInputSchema.parse(request.params.arguments); + const orgId = String(org_id); + span?.setAttribute("requested_org_id", orgId); + + // Mint (or reuse) an org-scoped bearer token for the requested org. No + // pre-validation against list_orgs: if the user can't access the org, the + // mint returns 401, which we surface as "org not accessible". + try { + await this.ensureOrgToken(orgId, recorder); + } catch (error) { + const message = (error as Error)?.message ?? ""; + if (message.includes("401")) { + return this.createErrorResponse( + `You do not have access to org "${orgId}", or it does not exist. Call list_orgs to see the orgs you can access.`, + "Switch org failed: org not accessible (401)", + ); + } + return this.createErrorResponse( + `Failed to switch to org "${orgId}". Please try again.`, + `Error switching org ${message}`, + ); + } + + await this.setActiveOrg(orgId); + // Data sources are org-specific; drop the cached set so the next lookup + // reflects the newly selected org. + this._sources = null; + span?.setAttribute("active_org_id", orgId); + + return this.createStructuredContentSuccessResponse( + { success: true, active_org_id: org_id }, + `Switched to org ${orgId}`, + ); + } + private _sources: { list: DataSource[]; map: Map; diff --git a/src/servers/tool-definitions.ts b/src/servers/tool-definitions.ts index f2c6f96..48f91f3 100644 --- a/src/servers/tool-definitions.ts +++ b/src/servers/tool-definitions.ts @@ -270,11 +270,31 @@ export const ListOrgsOutputSchema = z.object({ .string() .optional() .describe("Status of the Org (ACTIVE or IN_ACTIVE)."), + is_active: z + .boolean() + .describe( + "Whether this is the Org currently active for the session. Tool calls operate against the active Org.", + ), }), ) .describe("The list of Orgs the user can access."), }); +export const SwitchOrgInputSchema = z.object({ + org_id: z + .number() + .describe( + "The id of the Org to switch to. Use an `id` returned by `list_orgs`. Subsequent tool calls in this session operate against this Org.", + ), +}); + +export const SwitchOrgOutputSchema = z.object({ + success: z.boolean().describe("Whether the Org switch was successful."), + active_org_id: z + .number() + .describe("The id of the Org now active for the session."), +}); + export enum ToolName { // V1 Ping = "ping", @@ -289,6 +309,7 @@ export enum ToolName { GetSessionUpdates = "get_session_updates", CreateDashboard = "create_dashboard", ListOrgs = "list_orgs", + SwitchOrg = "switch_org", } export const toolDefinitionsV1 = [ @@ -420,7 +441,7 @@ export const toolDefinitionsV2 = [ { name: ToolName.ListOrgs, description: - "List the Orgs configured on the ThoughtSpot instance, including the ID, name, description, and status of each Org, along with the ID of the Org that is currently active for this session. Only available when authenticated via OAuth.", + "List the Orgs the authenticated user can access on the ThoughtSpot instance, including the ID, name, description, and status of each Org. The Org marked `is_active: true` is the one currently active for this session, which all tool calls operate against. Use this to tell the user which Org they are in. Only available when authenticated via OAuth.", inputSchema: z.toJSONSchema(ListOrgsInputSchema), outputSchema: z.toJSONSchema(ListOrgsOutputSchema), annotations: { @@ -430,4 +451,17 @@ export const toolDefinitionsV2 = [ openWorldHint: false, }, }, + { + name: ToolName.SwitchOrg, + description: + "Switch the active Org for this session. After switching, all subsequent tool calls (analysis sessions, answers, data sources, dashboards) operate against the selected Org. Pass an `org_id` returned by `list_orgs`. If you do not have access to the Org, the switch fails. The active Org is per-session and resets if the connection is re-established. Only available when authenticated via OAuth.", + inputSchema: z.toJSONSchema(SwitchOrgInputSchema), + outputSchema: z.toJSONSchema(SwitchOrgOutputSchema), + annotations: { + title: "Switch Org", + readOnlyHint: false, + destructiveHint: false, + openWorldHint: false, + }, + }, ]; diff --git a/src/storage-service/storage-service.ts b/src/storage-service/storage-service.ts index a57660c..48710e7 100644 --- a/src/storage-service/storage-service.ts +++ b/src/storage-service/storage-service.ts @@ -37,6 +37,41 @@ export class StorageServiceClient { return `https://internal/storage/${encodeURIComponent(conversationId)}/${operation}`; } + // Reserved pseudo-conversation id for the per-user active-org DO instance. Uses + // the same hash-namespaced addressing as conversations, but a distinct instance + // that holds only the active org (no message TTL), shared across the user's + // sessions. + private static readonly ACTIVE_ORG_ID = "__active_org__"; + + /** Read the user's active org id, or null if none has been set. */ + async getActiveOrg(): Promise { + const id = StorageServiceClient.ACTIVE_ORG_ID; + const response = await this.stubFor(id).fetch(this.url(id, "active-org"), { + method: "GET", + headers: this.headers(), + }); + if (!response.ok) { + const body = await response.text(); + throw new Error(`Failed to get active org (${response.status}): ${body}`); + } + const data = (await response.json()) as { activeOrgId: string | null }; + return data.activeOrgId ?? null; + } + + /** Persist the user's active org id (shared across their sessions). */ + async setActiveOrg(activeOrgId: string): Promise { + const id = StorageServiceClient.ACTIVE_ORG_ID; + const response = await this.stubFor(id).fetch(this.url(id, "active-org"), { + method: "POST", + headers: this.headers(), + body: JSON.stringify({ activeOrgId }), + }); + if (!response.ok) { + const body = await response.text(); + throw new Error(`Failed to set active org (${response.status}): ${body}`); + } + } + /** * Initialize a conversation. Must be called before appending messages. * Can also be called on an existing conversation that is already marked done, diff --git a/src/thoughtspot/thoughtspot-client.ts b/src/thoughtspot/thoughtspot-client.ts index cf88054..52d9815 100644 --- a/src/thoughtspot/thoughtspot-client.ts +++ b/src/thoughtspot/thoughtspot-client.ts @@ -1,23 +1,29 @@ import { - createBearerAuthenticationConfig, ThoughtSpotRestApi, + createBearerAuthenticationConfig, } from "@thoughtspot/rest-api-sdk"; import type { AgentConversation, RequestContext, ResponseContext, } from "@thoughtspot/rest-api-sdk"; -import YAML from "yaml"; +import { customAlphabet } from "nanoid"; import { of } from "rxjs"; +import YAML from "yaml"; import type { SessionInfo } from "./types"; -import { customAlphabet } from "nanoid"; /* * Inject custom handlers into the ThoughtSpot client */ +// Header used by ThoughtSpot to select which org a request operates against. +// The same access token works across all orgs the user belongs to; the active +// org is chosen per-request via this header. +const ORG_HEADER = "x-thoughtspot-orgs"; + export const getThoughtSpotClient = ( instanceUrl: string, bearerToken: string, + orgId?: string, ) => { const config = createBearerAuthenticationConfig(instanceUrl, () => Promise.resolve(bearerToken), @@ -29,6 +35,10 @@ export const getThoughtSpotClient = ( if (!headers || !headers["Accept-Language"]) { context.setHeaderParam("Accept-Language", "en-US"); } + // Scope every SDK call to the active org, if one is set. + if (orgId) { + context.setHeaderParam(ORG_HEADER, orgId); + } return of(context) as any; }, post: (context: ResponseContext) => { @@ -37,14 +47,47 @@ export const getThoughtSpotClient = ( }); const client = new ThoughtSpotRestApi(config); (client as any).instanceUrl = instanceUrl; - addExportUnsavedAnswerTML(client, instanceUrl, bearerToken); - addGetSessionInfo(client, instanceUrl, bearerToken); - addGetAnswerSession(client, instanceUrl, bearerToken); - addCreateAgentConversationWithAutoMode(client, instanceUrl, bearerToken); - addSendAgentConversationMessageStreaming(client, instanceUrl, bearerToken); + addExportUnsavedAnswerTML(client, instanceUrl, bearerToken, orgId); + addGetSessionInfo(client, instanceUrl, bearerToken, orgId); + addGetAnswerSession(client, instanceUrl, bearerToken, orgId); + addCreateAgentConversationWithAutoMode( + client, + instanceUrl, + bearerToken, + orgId, + ); + addSendAgentConversationMessageStreaming( + client, + instanceUrl, + bearerToken, + orgId, + ); + addGetRefreshedToken(client, instanceUrl); + addFetchOrgBearerToken(client, instanceUrl); return client; }; +/* + * Build the auth/content headers for the custom raw-fetch handlers below, + * including the org-scoping header when an active org is set. + */ +function buildHeaders( + token: string, + orgId?: string, + accept = "application/json", +): Record { + const headers: Record = { + "Content-Type": "application/json", + Accept: accept, + "user-agent": "ThoughtSpot-ts-client", + Authorization: `Bearer ${token}`, + }; + if (orgId) { + headers[ORG_HEADER] = orgId; + } + return headers; +} + const getAnswerTML = ` mutation GetUnsavedAnswerTML($session: BachSessionIdInput!, $exportDependencies: Boolean, $formatType: EDocFormatType, $exportPermissions: Boolean, $exportFqn: Boolean) { UnsavedAnswer_getTML( @@ -72,6 +115,7 @@ function addExportUnsavedAnswerTML( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).exportUnsavedAnswerTML = async ({ session_identifier, @@ -81,12 +125,7 @@ function addExportUnsavedAnswerTML( // make a graphql request to `ThoughtspotHost/prism endpoint. const response = await fetch(`${instanceUrl}${endpoint}`, { method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - "user-agent": "ThoughtSpot-ts-client", - Authorization: `Bearer ${token}`, - }, + headers: buildHeaders(token, orgId), body: JSON.stringify({ operationName: "GetUnsavedAnswerTML", query: getAnswerTML, @@ -112,18 +151,14 @@ async function addGetSessionInfo( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).getSessionInfo = async (): Promise => { const endpoint = "/prism/preauth/info"; // make a graphql request to `ThoughtspotHost/prism endpoint. const response = await fetch(`${instanceUrl}${endpoint}`, { method: "GET", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - "user-agent": "ThoughtSpot-ts-client", - Authorization: `Bearer ${token}`, - }, + headers: buildHeaders(token, orgId), }); const data: any = await response.json(); @@ -158,7 +193,12 @@ export interface AnswerSession { /* * Using custom handler because we don't have a public API for this */ -function addGetAnswerSession(client: any, instanceUrl: string, token: string) { +function addGetAnswerSession( + client: any, + instanceUrl: string, + token: string, + orgId?: string, +) { (client as any).getAnswerSession = async ({ session_identifier, generation_number, @@ -170,12 +210,7 @@ function addGetAnswerSession(client: any, instanceUrl: string, token: string) { const operationName = "Answer__updateTokens"; const fetchOptions = { method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - "user-agent": "ThoughtSpot-ts-client", - Authorization: `Bearer ${token}`, - }, + headers: buildHeaders(token, orgId), body: JSON.stringify({ operationName, query: getAnswerSessionQuery, @@ -211,6 +246,7 @@ function addCreateAgentConversationWithAutoMode( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).createAgentConversationWithAutoMode = async ({ dataSourceId, @@ -220,12 +256,7 @@ function addCreateAgentConversationWithAutoMode( const endpoint = "/conversation/v2/"; const fetchOptions = { method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - "user-agent": "ThoughtSpot-ts-client", - Authorization: `Bearer ${token}`, - }, + headers: buildHeaders(token, orgId), body: JSON.stringify({ context: dataSourceId ? { @@ -282,6 +313,7 @@ function addSendAgentConversationMessageStreaming( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).sendAgentConversationMessageStreaming = async ({ conversation_identifier, @@ -294,12 +326,7 @@ function addSendAgentConversationMessageStreaming( const endpoint = `/conversation/v2/${encodeURIComponent(conversation_identifier)}/query`; const fetchOptions = { method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "text/event-stream", - "user-agent": "ThoughtSpot-ts-client", - Authorization: `Bearer ${token}`, - }, + headers: buildHeaders(token, orgId, "text/event-stream"), body: JSON.stringify({ mode: "spotter", // TODO(Rifdhan) support deep analysis mode id: generateNanoID(), @@ -325,3 +352,104 @@ function addSendAgentConversationMessageStreaming( return response; }; } + +export interface RefreshedTokens { + accessToken: string; + refreshToken?: string; +} + +/* + * Fetches a fresh (cluster-wide) access token via the session gettoken endpoint + * with refresh=true, authenticated with the caller's current bearer token. + * Returns both tokens so they can be stored together. The returned token is not + * pinned to a single org; use fetchOrgBearerToken to mint an org-scoped one. + */ +function addGetRefreshedToken(client: any, instanceUrl: string) { + (client as any).getRefreshedToken = async ({ + bearerToken, + refreshToken, + }: { + bearerToken: string; + refreshToken?: string; + }): Promise => { + const endpoint = "/callosum/v1/session/v2/gettoken?refresh=true"; + const headers = buildHeaders(bearerToken); + if (refreshToken) { + headers["X-Refresh-Token"] = refreshToken; + } + const response = await fetch(`${instanceUrl}${endpoint}`, { + method: "GET", + headers, + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `getRefreshedToken failed with status ${response.status}: ${errorText}`, + ); + } + + const data = (await response.json()) as any; + const accessToken = data?.token ?? data?.data?.token; + if (!accessToken || typeof accessToken !== "string") { + throw new Error("getRefreshedToken: no token in response"); + } + const newRefreshToken = data?.refreshToken ?? data?.data?.refreshToken; + return { + accessToken, + refreshToken: + typeof newRefreshToken === "string" ? newRefreshToken : undefined, + }; + }; +} + +// Default validity for a minted org-scoped bearer token (30 days, in seconds), +// matching the validity the connector uses at login. +const ORG_TOKEN_VALIDITY_SEC = 30 * 24 * 60 * 60; + +/* + * Mints an ORG-SCOPED bearer token for the given org, authenticated with the + * caller's (cluster-wide) access token. Uses the Callosum v2 auth/token/fetch + * endpoint with org_identifier; the returned token is pinned to that org + * server-side. + * + * Note: the working path on these clusters is /callosum/v1/v2/auth/token/fetch + * (the /callosum/v2/... path 404s), and the token is nested under data.token. + */ +function addFetchOrgBearerToken(client: any, instanceUrl: string) { + (client as any).fetchOrgBearerToken = async ({ + accessToken, + orgId, + validityTimeInSec = ORG_TOKEN_VALIDITY_SEC, + }: { + accessToken: string; + orgId: string; + validityTimeInSec?: number; + }): Promise => { + const params = new URLSearchParams({ + validity_time_in_sec: String(validityTimeInSec), + org_identifier: orgId, + }); + const endpoint = `/callosum/v1/v2/auth/token/fetch?${params.toString()}`; + const response = await fetch(`${instanceUrl}${endpoint}`, { + method: "GET", + // Authenticate with the access token; no org header (the org is selected + // via org_identifier and pinned into the returned token). + headers: buildHeaders(accessToken), + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `fetchOrgBearerToken failed with status ${response.status}: ${errorText}`, + ); + } + + const data = (await response.json()) as any; + const token = data?.data?.token ?? data?.token; + if (!token || typeof token !== "string") { + throw new Error("fetchOrgBearerToken: no token in response"); + } + return token; + }; +} diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 2a4b19a..d0fe5b4 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -21,6 +21,7 @@ import { } from "../metrics/runtime/tool-metrics"; import { WithSpan, getActiveSpan } from "../metrics/tracing/tracing-utils"; import { processSendAgentConversationMessageStreamingResponse } from "../streaming-utils"; +import type { RefreshedTokens } from "./thoughtspot-client"; import type { Answer, DataSource, @@ -715,6 +716,7 @@ export class ThoughtSpotService { privileges: info.privileges, enableSpotterDataSourceDiscovery: info.configInfo?.enableSpotterDataSourceDiscovery, + orgsEnabled: info.configInfo?.orgsConfiguration?.enabled, }; } @@ -751,6 +753,38 @@ export class ThoughtSpotService { return results; } + /** + * Fetch a fresh (cluster-wide) access token via gettoken with refresh=true, + * authenticated with the given bearer token. + */ + @WithSpan("get-refreshed-token") + async getRefreshedToken( + bearerToken: string, + refreshToken?: string, + ): Promise { + return (await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getRefreshedToken, + () => + (this.client as any).getRefreshedToken({ bearerToken, refreshToken }), + )) as RefreshedTokens; + } + + /** + * Mint an org-scoped bearer token for `orgId`, authenticated with the given + * (cluster-wide) access token. + */ + @WithSpan("fetch-org-bearer-token") + async fetchOrgBearerToken( + accessToken: string, + orgId: string, + ): Promise { + getActiveSpan()?.setAttribute("org_id", orgId); + return (await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.fetchOrgBearerToken, + () => (this.client as any).fetchOrgBearerToken({ accessToken, orgId }), + )) as string; + } + /** * Search worksheets by term */ diff --git a/src/thoughtspot/types.ts b/src/thoughtspot/types.ts index cca4f1d..2e487be 100644 --- a/src/thoughtspot/types.ts +++ b/src/thoughtspot/types.ts @@ -35,6 +35,9 @@ export interface SessionInfo { currentOrgId: string; privileges: any; enableSpotterDataSourceDiscovery?: boolean; + // Whether Orgs are enabled on this cluster (configInfo.orgsConfiguration.enabled). + // Gates the org tools (list_orgs/switch_org). + orgsEnabled?: boolean; } export interface BaseMessage { diff --git a/src/utils.ts b/src/utils.ts index 0dd85d5..b37eef3 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -4,6 +4,22 @@ import { getActiveSpan } from "./metrics/tracing/tracing-utils"; export type Props = { accessToken: string; + /** + * Refresh token returned alongside the global access token by + * callosum/v1/session/v2/gettoken. Stored for future token refresh; not yet + * used to refresh automatically. + */ + refreshToken?: string; + /** + * When the global access token was issued (epoch millis), from gettoken's + * `tokenCreatedTime`. Stored as-is. + */ + tokenCreatedTime?: number; + /** + * When the global access token expires (epoch millis), from gettoken's + * `tokenExpiryDuration` (an absolute timestamp despite the name). Stored as-is. + */ + tokenExpiryDuration?: number; instanceUrl: string; clientName: { clientId: string; diff --git a/static/oauth-callback.js b/static/oauth-callback.js index ee65d71..ef99d44 100644 --- a/static/oauth-callback.js +++ b/static/oauth-callback.js @@ -18,7 +18,7 @@ if (!base.pathname.endsWith('/')) { base.pathname += '/'; } - const tokenUrl = new URL('callosum/v1/v2/auth/token/fetch?validity_time_in_sec=2592000', base.toString()); + const tokenUrl = new URL('callosum/v1/session/v2/gettoken?refresh=true', base.toString()); document.getElementById('status').textContent = 'Retrieving authentication token...'; @@ -130,7 +130,20 @@ } } - const data = await response.json(); + const raw = await response.json(); + // gettoken returns the access token at top-level `token`, plus refreshToken + // and the token's created/expiry timestamps. Normalize to the + // { data: { token, ... } } shape that /store-token expects (same shape the + // manual-paste path produces), carrying all fields through as-is. + const src = (raw && raw.data) ? raw.data : raw; + const data = { + data: { + token: src.token, + refreshToken: src.refreshToken, + tokenCreatedTime: src.tokenCreatedTime, + tokenExpiryDuration: src.tokenExpiryDuration, + }, + }; document.getElementById('status').textContent = 'Authentication successful. Securing your session...'; // Send the token to the server diff --git a/test/servers/mcp-server-base.spec.ts b/test/servers/mcp-server-base.spec.ts index 5c40760..ac5bce0 100644 --- a/test/servers/mcp-server-base.spec.ts +++ b/test/servers/mcp-server-base.spec.ts @@ -73,6 +73,18 @@ class TestMCPServer extends MCPServer { public async testGetStorageService() { return this.getStorageService(); } + + public async testGetStorageKeyHash() { + return this.getStorageKeyHash(); + } + + public testIsOrgsEnabled() { + return this.isOrgsEnabled(); + } + + public setSessionInfo(info: any) { + this.sessionInfo = info; + } } describe("MCP Server Base", () => { @@ -418,7 +430,7 @@ describe("MCP Server Base", () => { env: mockEnv, }); await expect(serverWithNoToken.testGetStorageService()).rejects.toThrow( - "Access token is required to use Storage Service", + "A token is required to derive the storage key", ); }); @@ -428,7 +440,7 @@ describe("MCP Server Base", () => { env: mockEnv, }); await expect(serverWithNoToken.testGetStorageService()).rejects.toThrow( - "Access token is required to use Storage Service", + "A token is required to derive the storage key", ); }); @@ -508,6 +520,86 @@ describe("MCP Server Base", () => { }); }); + describe("getStorageKeyHash", () => { + async function sha256Base64Url(token: string): Promise { + const buf = await crypto.subtle.digest( + "SHA-256", + new TextEncoder().encode(token), + ); + return Buffer.from(new Uint8Array(buf)).toString("base64url"); + } + + it("prefers the refresh token when present (stable across access-token rotation)", async () => { + const s = new TestMCPServer({ + props: { + ...mockProps, + accessToken: "access-A", + refreshToken: "refresh-X", + }, + env: mockEnv, + }); + expect(await s.testGetStorageKeyHash()).toBe( + await sha256Base64Url("refresh-X"), + ); + }); + + it("is stable when the access token rotates but refresh token is unchanged", async () => { + const s1 = new TestMCPServer({ + props: { ...mockProps, accessToken: "access-A", refreshToken: "r" }, + env: mockEnv, + }); + const s2 = new TestMCPServer({ + props: { ...mockProps, accessToken: "access-B", refreshToken: "r" }, + env: mockEnv, + }); + expect(await s1.testGetStorageKeyHash()).toBe( + await s2.testGetStorageKeyHash(), + ); + }); + + it("falls back to the access token when no refresh token (bearer/token auth)", async () => { + const s = new TestMCPServer({ + props: { + ...mockProps, + accessToken: "access-only", + refreshToken: undefined, + }, + env: mockEnv, + }); + expect(await s.testGetStorageKeyHash()).toBe( + await sha256Base64Url("access-only"), + ); + }); + + it("throws when neither token is present", async () => { + const s = new TestMCPServer({ + props: { ...mockProps, accessToken: "", refreshToken: undefined }, + env: mockEnv, + }); + await expect(s.testGetStorageKeyHash()).rejects.toThrow( + "A token is required to derive the storage key", + ); + }); + }); + + describe("isOrgsEnabled", () => { + it("returns false when session info is not initialized (fail closed)", () => { + expect(server.testIsOrgsEnabled()).toBe(false); + }); + + it("returns true when orgsEnabled is true", () => { + server.setSessionInfo({ orgsEnabled: true }); + expect(server.testIsOrgsEnabled()).toBe(true); + }); + + it("returns false when orgsEnabled is false or absent", () => { + server.setSessionInfo({ orgsEnabled: false }); + expect(server.testIsOrgsEnabled()).toBe(false); + server.setSessionInfo({}); + expect(server.testIsOrgsEnabled()).toBe(false); + }); + }); + describe("Server Initialization", () => { it("should be defined after construction", () => { expect(server).toBeDefined(); diff --git a/test/servers/mcp-server-orgs.spec.ts b/test/servers/mcp-server-orgs.spec.ts new file mode 100644 index 0000000..39ace1e --- /dev/null +++ b/test/servers/mcp-server-orgs.spec.ts @@ -0,0 +1,261 @@ +import { connect } from "mcp-testing-kit"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { MCPServer } from "../../src/servers/mcp-server"; +import * as thoughtspotClient from "../../src/thoughtspot/thoughtspot-client"; + +vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ + MixpanelTracker: vi.fn().mockImplementation(() => ({ track: vi.fn() })), +})); + +/** + * Tests for the org tools (list_orgs / switch_org) and their supporting + * machinery: the OAuth + orgs-enabled gate, the shared active-org store, and + * org-scoped token minting. + */ + +// A fake CONVERSATION_STORAGE_OBJECT namespace that emulates the active-org DO. +// Stores active-org values in a Map keyed by the DO instance name, so reads and +// writes from any server instance sharing the same storage-key hash see the same +// value (mirroring the real shared store). +function makeStorageNamespace(store: Map) { + return { + idFromName: (name: string) => ({ name }), + get: (id: { name: string }) => ({ + fetch: async (url: string, init?: RequestInit) => { + const op = new URL(url).pathname.split("/").pop(); + if (op === "active-org" && (init?.method ?? "GET") === "GET") { + return Response.json({ activeOrgId: store.get(id.name) ?? null }); + } + if (op === "active-org" && init?.method === "POST") { + const body = JSON.parse(String(init?.body)) as { + activeOrgId: string; + }; + store.set(id.name, body.activeOrgId); + return Response.json({ ok: true }); + } + return new Response("Not Found", { status: 404 }); + }, + }), + }; +} + +type SessionInfoOverrides = { + orgsEnabled?: boolean; + currentOrgId?: string; +}; + +function makeClientMock(opts: { + session?: SessionInfoOverrides; + orgs?: Array<{ id: number; name: string; status: string }>; + fetchOrgBearerToken?: ReturnType; +}) { + const orgsConfiguration = + opts.session?.orgsEnabled === undefined + ? undefined + : { enabled: opts.session.orgsEnabled }; + return { + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "10.13.0.cl-110", + userGUID: "test-user-123", + userName: "test-user", + currentOrgId: opts.session?.currentOrgId ?? "0", + privileges: [], + configInfo: { + mixpanelConfig: { + devSdkKey: "k", + prodSdkKey: "k", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + enableSpotterDataSourceDiscovery: false, + orgsConfiguration, + }, + }), + searchOrgs: vi.fn().mockResolvedValue( + opts.orgs ?? [ + { id: 0, name: "Primary", status: "ACTIVE", description: "Primary" }, + { id: 101, name: "DataPlatform", status: "ACTIVE" }, + ], + ), + fetchOrgBearerToken: + opts.fetchOrgBearerToken ?? vi.fn().mockResolvedValue("org-scoped-token"), + instanceUrl: "https://test.thoughtspot.cloud", + } as any; +} + +function makeServer(opts: { + authMode?: string; + session?: SessionInfoOverrides; + orgs?: Array<{ id: number; name: string; status: string }>; + store?: Map; + fetchOrgBearerToken?: ReturnType; +}) { + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue( + makeClientMock(opts), + ); + const store = opts.store ?? new Map(); + const env = { + CONVERSATION_STORAGE_OBJECT: makeStorageNamespace(store), + } as any; + const props = { + instanceUrl: "https://test.thoughtspot.cloud", + accessToken: "global-token", + refreshToken: "refresh-token", + authMode: opts.authMode, + apiVersion: "latest", + clientName: { + clientId: "c", + clientName: "c", + registrationDate: 0, + }, + }; + return { server: new MCPServer({ props, env }), store }; +} + +describe("MCP Server org tools", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("tool visibility gate (OAuth AND orgs enabled)", () => { + it("lists org tools when OAuth and orgs enabled", async () => { + const { server } = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true }, + }); + await server.init(); + const { listTools } = connect(server); + const names = (await listTools()).tools?.map((t) => t.name) ?? []; + expect(names).toContain("list_orgs"); + expect(names).toContain("switch_org"); + }); + + it("hides org tools when orgs are not enabled on the cluster", async () => { + const { server } = makeServer({ + authMode: "oauth", + session: { orgsEnabled: false }, + }); + await server.init(); + const { listTools } = connect(server); + const names = (await listTools()).tools?.map((t) => t.name) ?? []; + expect(names).not.toContain("list_orgs"); + expect(names).not.toContain("switch_org"); + }); + + it("hides org tools for non-OAuth (bearer) connections even if orgs enabled", async () => { + const { server } = makeServer({ + authMode: "bearer", + session: { orgsEnabled: true }, + }); + await server.init(); + const { listTools } = connect(server); + const names = (await listTools()).tools?.map((t) => t.name) ?? []; + expect(names).not.toContain("list_orgs"); + expect(names).not.toContain("switch_org"); + }); + + it("fails closed: hides org tools when orgs-enabled flag is absent", async () => { + const { server } = makeServer({ + authMode: "oauth", + session: {}, // orgsConfiguration undefined + }); + await server.init(); + const { listTools } = connect(server); + const names = (await listTools()).tools?.map((t) => t.name) ?? []; + expect(names).not.toContain("list_orgs"); + expect(names).not.toContain("switch_org"); + }); + }); + + describe("list_orgs", () => { + it("returns ACTIVE orgs and marks the current org active when none switched", async () => { + const { server } = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true, currentOrgId: "0" }, + }); + await server.init(); + const { callTool } = connect(server); + const res = await callTool("list_orgs", {}); + const data = JSON.parse(res.content[0].text); + expect(data.orgs.map((o: any) => o.id)).toEqual([0, 101]); + // No explicit switch has happened, so nothing is stored as active and no + // org is marked is_active (calls use the cluster-resolved default org). + expect(data.orgs.every((o: any) => o.is_active === false)).toBe(true); + }); + + it("rejects direct invocation when org tools are unavailable", async () => { + const { server } = makeServer({ + authMode: "bearer", + session: { orgsEnabled: true }, + }); + await server.init(); + const { callTool } = connect(server); + const res = await callTool("list_orgs", {}); + expect(res.isError).toBe(true); + }); + }); + + describe("switch_org", () => { + it("mints an org token and persists the active org to the shared store", async () => { + const { server, store } = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true, currentOrgId: "0" }, + }); + await server.init(); + + const switchRes = await connect(server).callTool("switch_org", { + org_id: 101, + }); + const switchData = JSON.parse(switchRes.content[0].text); + expect(switchData.success).toBe(true); + expect(switchData.active_org_id).toBe(101); + // Persisted to the shared store. + expect([...store.values()]).toContain("101"); + }); + + it("returns 'not accessible' when minting the org token 401s", async () => { + const { server } = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true }, + fetchOrgBearerToken: vi + .fn() + .mockRejectedValue( + new Error("fetchOrgBearerToken failed with status 401: nope"), + ), + }); + await server.init(); + const { callTool } = connect(server); + const res = await callTool("switch_org", { org_id: 999 }); + expect(res.isError).toBe(true); + expect(res.content[0].text).toMatch(/do not have access/i); + }); + }); + + describe("shared active-org store persists across server instances", () => { + it("a switch in one instance is visible to another instance with the same token", async () => { + const store = new Map(); + const a = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true }, + store, + }); + await a.server.init(); + await connect(a.server).callTool("switch_org", { org_id: 101 }); + + // Second server instance (e.g. a different MCP session/DO) sharing the + // same store + token. + const b = makeServer({ + authMode: "oauth", + session: { orgsEnabled: true }, + store, + }); + await b.server.init(); + const listRes = await connect(b.server).callTool("list_orgs", {}); + const listData = JSON.parse(listRes.content[0].text); + expect(listData.orgs.find((o: any) => o.is_active).id).toBe(101); + }); + }); +}); diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index 614bd57..7703d5f 100644 --- a/test/thoughtspot/thoughtspot-service.spec.ts +++ b/test/thoughtspot/thoughtspot-service.spec.ts @@ -1852,4 +1852,57 @@ describe("thoughtspot-service", () => { ]); }); }); + + describe("searchOrgs", () => { + it("returns ACTIVE orgs only, mapping id/name/description/status", async () => { + mockClient.searchOrgs = vi.fn().mockResolvedValue([ + { id: 0, name: "Primary", status: "ACTIVE", description: "P" }, + { id: 5, name: "Inactive", status: "IN_ACTIVE" }, + { id: 9, name: "Data", status: "ACTIVE" }, + ]); + const service = new ThoughtSpotService(mockClient); + const orgs = await service.searchOrgs(); + expect(mockClient.searchOrgs).toHaveBeenCalledWith({ status: "ACTIVE" }); + expect(orgs).toEqual([ + { id: 0, name: "Primary", description: "P", status: "ACTIVE" }, + { id: 9, name: "Data", description: undefined, status: "ACTIVE" }, + ]); + }); + + it("tolerates a nullish upstream response", async () => { + mockClient.searchOrgs = vi.fn().mockResolvedValue(undefined); + const service = new ThoughtSpotService(mockClient); + await expect(service.searchOrgs()).resolves.toEqual([]); + }); + }); + + describe("fetchOrgBearerToken", () => { + it("delegates to the client with the access token and org id", async () => { + mockClient.fetchOrgBearerToken = vi + .fn() + .mockResolvedValue("org-token-xyz"); + const service = new ThoughtSpotService(mockClient); + const token = await service.fetchOrgBearerToken("global-tok", "101"); + expect(mockClient.fetchOrgBearerToken).toHaveBeenCalledWith({ + accessToken: "global-tok", + orgId: "101", + }); + expect(token).toBe("org-token-xyz"); + }); + }); + + describe("getRefreshedToken", () => { + it("passes the bearer and refresh tokens to the client", async () => { + mockClient.getRefreshedToken = vi + .fn() + .mockResolvedValue({ accessToken: "new-a", refreshToken: "same-r" }); + const service = new ThoughtSpotService(mockClient); + const res = await service.getRefreshedToken("bear", "ref"); + expect(mockClient.getRefreshedToken).toHaveBeenCalledWith({ + bearerToken: "bear", + refreshToken: "ref", + }); + expect(res).toEqual({ accessToken: "new-a", refreshToken: "same-r" }); + }); + }); });