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/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/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..f0a637c 100644 --- a/src/metrics/runtime/tool-metrics.ts +++ b/src/metrics/runtime/tool-metrics.ts @@ -21,6 +21,9 @@ export const UPSTREAM_OPERATION_NAMES = { "send_agent_conversation_message_streaming", 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 334b551..859b1ab 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -26,6 +26,7 @@ import { GetRelevantQuestionsSchema, GetSessionUpdatesInputSchema, SendSessionMessageInputSchema, + SwitchOrgInputSchema, ToolName, } from "./tool-definitions"; import { @@ -35,14 +36,137 @@ 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"; } + /** + * 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"; + } + + /** + * 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) { @@ -136,6 +260,15 @@ export class MCPServer extends BaseMCPServer { ); } + // 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 }; } @@ -192,7 +325,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 +387,30 @@ export class MCPServer extends BaseMCPServer { return this.callCreateDashboard(request, recorder); } + case ToolName.ListOrgs: { + // 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 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}`); } @@ -612,6 +769,72 @@ 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); + + // 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.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 8db2e91..48f91f3 100644 --- a/src/servers/tool-definitions.ts +++ b/src/servers/tool-definitions.ts @@ -257,6 +257,44 @@ 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)."), + 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", @@ -270,6 +308,8 @@ export enum ToolName { SendSessionMessage = "send_session_message", GetSessionUpdates = "get_session_updates", CreateDashboard = "create_dashboard", + ListOrgs = "list_orgs", + SwitchOrg = "switch_org", } export const toolDefinitionsV1 = [ @@ -398,4 +438,30 @@ export const toolDefinitionsV2 = [ openWorldHint: false, }, }, + { + name: ToolName.ListOrgs, + description: + "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: { + title: "List Orgs", + readOnlyHint: true, + destructiveHint: false, + 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 974cb17..d0fe5b4 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -21,11 +21,13 @@ 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, DataSourceSuggestion, Message, + Org, SessionInfo, } from "./types"; @@ -714,9 +716,75 @@ export class ThoughtSpotService { privileges: info.privileges, enableSpotterDataSourceDiscovery: info.configInfo?.enableSpotterDataSourceDiscovery, + orgsEnabled: info.configInfo?.orgsConfiguration?.enabled, }; } + /** + * 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; + } + + /** + * 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 b4d6a60..2e487be 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; @@ -28,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 b7c525d..b37eef3 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1,9 +1,25 @@ 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 = { 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; @@ -13,6 +29,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/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/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 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" }); + }); + }); });