From 4b718e1d96f56e0d40d640a9415f0f9b4313800a Mon Sep 17 00:00:00 2001 From: Rohit Gupta Date: Thu, 18 Jun 2026 23:08:38 +0530 Subject: [PATCH] Add multi-org support: list_orgs and switch_org MCP tools Add two MCP tools that let a connected user discover the orgs they belong to and switch the active org mid-session, without re-authenticating. - list_orgs: lists the user's orgs (via session/orgs), marking the active one. - switch_org: sets the active org; mints and caches an org-scoped bearer token for it and uses it for that org's calls. Token model: - Global cluster token (props.accessToken) is used to list orgs and to mint per-org tokens. - Per-org tokens are minted on demand via auth/token/fetch with org_identifier and cached in-memory, keyed by org. - An org-scoped token for the session's current org is minted on connect (postInit) so an active org token always exists. - Requests carry the x-thoughtspot-orgs header for the active org. - Login (oauth-callback) fetches the token via gettoken?refresh=true. --- src/metrics/runtime/tool-metrics.ts | 3 + src/servers/mcp-server-base.ts | 61 ++++++- src/servers/mcp-server.ts | 219 ++++++++++++++++++++++- src/servers/tool-definitions.ts | 70 ++++++++ src/thoughtspot/thoughtspot-client.ts | 236 +++++++++++++++++++++---- src/thoughtspot/thoughtspot-service.ts | 43 +++++ src/thoughtspot/types.ts | 4 + static/oauth-callback.js | 11 +- 8 files changed, 603 insertions(+), 44 deletions(-) diff --git a/src/metrics/runtime/tool-metrics.ts b/src/metrics/runtime/tool-metrics.ts index c86b7ee..9282601 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", + listOrgs: "list_orgs", + fetchOrgBearerToken: "fetch_org_bearer_token", + getRefreshedToken: "get_refreshed_token", } as const; export type UpstreamOperation = diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 4df14f0..65b72df 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -207,6 +207,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 OAuth 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 listing and + * 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 +258,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 +484,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..ee2a572 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -13,6 +13,7 @@ import { } from "../metrics/runtime/metrics-recorder"; import type { ToolMetricApiSurface } from "../metrics/runtime/tool-metrics"; import { WithSpan } from "../metrics/tracing/tracing-utils"; +import type { Org } from "../thoughtspot/thoughtspot-client"; import type { DataSource } from "../thoughtspot/thoughtspot-service"; import type { Answer, StreamingMessagesState } from "../thoughtspot/types"; import { McpServerError } from "../utils"; @@ -26,6 +27,7 @@ import { GetRelevantQuestionsSchema, GetSessionUpdatesInputSchema, SendSessionMessageInputSchema, + SwitchOrgInputSchema, ToolName, } from "./tool-definitions"; import { @@ -35,10 +37,91 @@ import { } from "./version-registry"; export class MCPServer extends BaseMCPServer { + // Active org for this session, held in-memory on the Durable Object instance. + // Resets if the DO is evicted/restarted; the user can re-select via switch_org. + private activeOrgId: string | undefined; + + // Org-scoped bearer tokens, keyed by org id. Minted on demand from the + // global cluster token (props.accessToken) and reused for that org's calls. + private orgBearerTokens = new Map(); + constructor(ctx: Context) { super(ctx, "ThoughtSpot", "2.0.0"); } + protected getActiveOrgId(): string | undefined { + return this.activeOrgId; + } + + /** + * When an org is active and we hold a bearer token for it, use that token. + * Otherwise fall back to the session's default access token. + */ + protected getActiveBearerToken(): string { + if (this.activeOrgId) { + const orgToken = this.orgBearerTokens.get(this.activeOrgId); + if (orgToken) { + return orgToken; + } + } + return this.ctx.props.accessToken; + } + + /** + * The global (cluster-level) access token used to list orgs and to mint + * per-org tokens. This is the token minted at login by the browser's + * gettoken call and stored as props.accessToken (via /store-token); the + * server uses it directly. It is NOT pinned to a single org. + */ + private resolveAccessToken(_recorder: MetricsRecorder): string { + return this.ctx.props.accessToken; + } + + /** + * Ensure we hold an org-scoped token for `orgId`: reuse the cached one, or + * mint a new one from the global token and cache it. Returns the org token. + */ + private async ensureOrgToken( + orgId: string, + recorder?: MetricsRecorder, + ): Promise { + const existing = this.orgBearerTokens.get(orgId); + if (existing) { + return existing; + } + const orgToken = await this.getThoughtSpotServiceWithToken( + this.ctx.props.accessToken, + undefined, + recorder, + ).fetchOrgBearerToken(this.ctx.props.accessToken, orgId); + this.orgBearerTokens.set(orgId, orgToken); + return orgToken; + } + + /** + * On connection, mint an org-scoped token for the session's current org so + * there is always an active org token available (matching the original + * single-org behavior). Best-effort: failures don't break the connection. + */ + protected async postInit(): Promise { + const currentOrgId = + this.sessionInfo?.currentOrgId !== undefined + ? String(this.sessionInfo.currentOrgId) + : undefined; + if (!currentOrgId) { + return; + } + try { + await this.ensureOrgToken(currentOrgId); + this.activeOrgId = currentOrgId; + } catch (error) { + console.error( + `Failed to mint initial org token for org ${currentOrgId}:`, + error, + ); + } + } + protected getToolMetricApiSurface(): ToolMetricApiSurface { return "mcp"; } @@ -192,7 +275,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 +337,14 @@ export class MCPServer extends BaseMCPServer { return this.callCreateDashboard(request, recorder); } + case ToolName.ListOrgs: { + return this.callListOrgs(recorder); + } + + case ToolName.SwitchOrg: { + return this.callSwitchOrg(request, recorder); + } + default: throw new Error(`Unknown tool: ${name}`); } @@ -579,6 +670,132 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; ); } + @WithSpan("call-list-orgs") + async callListOrgs(recorder: MetricsRecorder) { + if (!this.ctx.props.accessToken || !this.ctx.props.instanceUrl) { + return this.createErrorResponse( + "Access token or instance URL not valid", + "List orgs failed", + ); + } + + // Use the global cluster token (from login) to list orgs. + const accessToken = this.resolveAccessToken(recorder); + + let orgs: Org[]; + try { + orgs = await this.getThoughtSpotServiceWithToken( + accessToken, + undefined, + recorder, + ).listOrgs(); + } catch (error) { + return this.createErrorResponse( + "Failed to list orgs. Your account may not have orgs enabled, or your authentication may have expired.", + `Error listing orgs ${(error as Error)?.message}`, + ); + } + + // Determine the currently active org. If the user has explicitly switched, + // use that; otherwise fall back to the org the session resolves to. + let activeOrgId = this.getActiveOrgId(); + if (!activeOrgId) { + try { + const sessionInfo = + await this.getThoughtSpotService(recorder).getSessionInfo(); + if (sessionInfo.currentOrgId !== undefined) { + activeOrgId = String(sessionInfo.currentOrgId); + } + } catch { + // Best-effort: if we can't resolve the current org, leave all inactive. + } + } + + return this.createStructuredContentSuccessResponse( + { + orgs: orgs.map((org) => ({ + id: org.id, + name: org.name, + description: org.description, + is_active: activeOrgId !== undefined && 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); + span?.setAttribute("requested_org_id", org_id); + + if (!this.ctx.props.accessToken || !this.ctx.props.instanceUrl) { + return this.createErrorResponse( + "Access token or instance URL not valid", + "Switch org failed", + ); + } + + // The global cluster token — used both to verify org access and to mint + // the org-scoped bearer token. + const accessToken = this.resolveAccessToken(recorder); + + // Validate the requested org against the orgs the user can actually access. + // This prevents switching to an org the user isn't a member of (which would + // otherwise surface as opaque 401/INACCESSIBLE_ORG errors on later calls). + let orgs: Org[]; + try { + orgs = await this.getThoughtSpotServiceWithToken( + accessToken, + undefined, + recorder, + ).listOrgs(); + } catch (error) { + return this.createErrorResponse( + "Failed to verify org access. Your authentication may have expired.", + `Error listing orgs while switching ${(error as Error)?.message}`, + ); + } + + const target = orgs.find((org) => org.id === org_id); + if (!target) { + return this.createErrorResponse( + `Org "${org_id}" is not in the list of orgs you have access to. Call list_orgs to see available orgs.`, + "Switch org failed: org not accessible", + ); + } + + // Ensure we hold an org-scoped bearer token for this org (reuse cached or + // mint from the global token). + try { + await this.ensureOrgToken(target.id, recorder); + } catch (error) { + return this.createErrorResponse( + `Failed to obtain an access token for org "${target.name}". Please try again.`, + `Error minting org bearer token ${(error as Error)?.message}`, + ); + } + + this.activeOrgId = target.id; + // 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", target.id); + + return this.createStructuredContentSuccessResponse( + { + success: true, + active_org_id: target.id, + active_org_name: target.name, + }, + `Switched to org ${target.name}`, + ); + } + @WithSpan("call-get-data-source-suggestions") async callGetDataSourceSuggestions( request: z.infer, diff --git a/src/servers/tool-definitions.ts b/src/servers/tool-definitions.ts index 8db2e91..0608a16 100644 --- a/src/servers/tool-definitions.ts +++ b/src/servers/tool-definitions.ts @@ -257,6 +257,48 @@ export const CreateDashboardOutputSchema = z.object({ ), }); +export const ListOrgsInputSchema = z.object({}); + +export const OrgSchema = z.object({ + id: z.string().describe("The unique identifier of the org."), + name: z.string().describe("The name of the org."), + description: z + .string() + .optional() + .describe("The description of the org, if any."), + is_active: z + .boolean() + .describe( + "Whether this org is the one currently active for the session. Tool calls operate against the active org.", + ), +}); + +export const ListOrgsOutputSchema = z.object({ + orgs: z + .array(OrgSchema) + .describe( + "The list of orgs the authenticated user is a member of. Use an org `id` with `switch_org` to make subsequent tool calls operate against that org.", + ), +}); + +export const SwitchOrgInputSchema = z.object({ + org_id: z + .string() + .describe( + "The id of the org to switch to. Must be one of the org ids returned by `list_orgs`. Subsequent tool calls in this session will operate against this org.", + ), +}); + +export const SwitchOrgOutputSchema = z.object({ + success: z.boolean().describe("Whether the org switch was successful."), + active_org_id: z + .string() + .describe("The id of the org now active for the session."), + active_org_name: z + .string() + .describe("The name of the org now active for the session."), +}); + export enum ToolName { // V1 Ping = "ping", @@ -270,6 +312,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 +442,30 @@ export const toolDefinitionsV2 = [ openWorldHint: false, }, }, + { + name: ToolName.ListOrgs, + description: + "List the orgs (tenants) the authenticated user is a member of. Each org has an `id` and `name`. The org marked `is_active: true` is the one currently used for tool calls. Use the returned `id` with `switch_org` to operate against a different org. Call this when the user asks which orgs they have access to, or before switching orgs.", + 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`. Note: the active org is per-session and resets if the connection is re-established.", + inputSchema: z.toJSONSchema(SwitchOrgInputSchema), + outputSchema: z.toJSONSchema(SwitchOrgOutputSchema), + annotations: { + title: "Switch Org", + readOnlyHint: false, + destructiveHint: false, + openWorldHint: false, + }, + }, ]; diff --git a/src/thoughtspot/thoughtspot-client.ts b/src/thoughtspot/thoughtspot-client.ts index cf88054..06b0771 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 bearer/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,48 @@ 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, + ); + addListOrgs(client, instanceUrl, bearerToken); + addFetchOrgBearerToken(client, instanceUrl); + addGetRefreshedToken(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 +116,7 @@ function addExportUnsavedAnswerTML( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).exportUnsavedAnswerTML = async ({ session_identifier, @@ -81,12 +126,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 +152,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 +194,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 +211,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 +247,7 @@ function addCreateAgentConversationWithAutoMode( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).createAgentConversationWithAutoMode = async ({ dataSourceId, @@ -220,12 +257,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 +314,7 @@ function addSendAgentConversationMessageStreaming( client: any, instanceUrl: string, token: string, + orgId?: string, ) { (client as any).sendAgentConversationMessageStreaming = async ({ conversation_identifier, @@ -294,12 +327,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 +353,133 @@ function addSendAgentConversationMessageStreaming( return response; }; } + +export interface Org { + id: string; + name: string; + description?: string; +} + +/* + * Lists the orgs the authenticated user is a member of, via the session orgs + * endpoint. We avoid the orgs/search REST endpoint because it is not permitted + * for regular users on these clusters (403 "Operation is not allowed"). This + * endpoint returns { orgs: [{orgId, orgName, description, isActive}], currentOrgId }. + */ +function addListOrgs(client: any, instanceUrl: string, token: string) { + (client as any).listOrgs = async (): Promise => { + const endpoint = "/callosum/v1/session/orgs?batchsize=-1&offset=-1"; + const response = await fetch(`${instanceUrl}${endpoint}`, { + method: "GET", + headers: buildHeaders(token), + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `listOrgs failed with status ${response.status}: ${errorText}`, + ); + } + + const data = (await response.json()) as any; + const orgs: any[] = Array.isArray(data?.orgs) ? data.orgs : []; + return orgs.map((org) => ({ + id: String(org.orgId ?? org.id), + name: org.orgName ?? org.name ?? String(org.orgId ?? org.id), + description: org.description || undefined, + })); + }; +} + +export interface RefreshedTokens { + accessToken: string; + refreshToken?: string; +} + +/* + * Fetches a fresh access token (and refresh token, if available) via the + * session gettoken endpoint with refresh=true. Authenticated with the caller's + * current bearer token (same cross-domain mechanism as the other handlers). + * Returns both tokens so they can be stored together. + */ +function addGetRefreshedToken(client: any, instanceUrl: string) { + (client as any).getRefreshedToken = async ({ + bearerToken, + }: { + bearerToken: string; + }): Promise => { + const endpoint = "/callosum/v1/session/v2/gettoken?refresh=true"; + const response = await fetch(`${instanceUrl}${endpoint}`, { + method: "GET", + headers: buildHeaders(bearerToken), + }); + + 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 refreshToken = data?.refreshToken ?? data?.data?.refreshToken; + return { + accessToken, + refreshToken: typeof refreshToken === "string" ? refreshToken : undefined, + }; + }; +} + +// Default validity for a minted org-scoped bearer token (30 days, in seconds), +// matching the validity the connector uses elsewhere. +const ORG_TOKEN_VALIDITY_SEC = 30 * 24 * 60 * 60; + +/* + * Mints an ORG-SCOPED bearer token for the given org, authenticated with the + * caller's access token. Uses the Callosum v2 auth/token/fetch endpoint with + * org_identifier. The returned token is pinned to that org server-side. + * `accessToken` here is the token fetched from session/info. + */ +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, + }); + // 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. + 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..7488160 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 { Org, RefreshedTokens } from "./thoughtspot-client"; import type { Answer, DataSource, @@ -714,9 +715,51 @@ export class ThoughtSpotService { privileges: info.privileges, enableSpotterDataSourceDiscovery: info.configInfo?.enableSpotterDataSourceDiscovery, + accessToken: info.accessToken, }; } + /** + * List the orgs the authenticated user is a member of. + */ + @WithSpan("list-orgs") + async listOrgs(): Promise { + const orgs = (await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.listOrgs, + () => (this.client as any).listOrgs(), + )) as Org[]; + getActiveSpan()?.setAttribute("total_orgs", orgs.length); + return orgs; + } + + /** + * Fetch a fresh access token (and refresh token, if any) via gettoken + * with refresh=true, authenticated with the given bearer token. + */ + @WithSpan("get-refreshed-token") + async getRefreshedToken(bearerToken: string): Promise { + return (await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getRefreshedToken, + () => (this.client as any).getRefreshedToken({ bearerToken }), + )) as RefreshedTokens; + } + + /** + * Mint an org-scoped bearer token for `orgId`, authenticated with the given + * access token (typically fetched from session/info). + */ + @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..aae7c83 100644 --- a/src/thoughtspot/types.ts +++ b/src/thoughtspot/types.ts @@ -28,6 +28,10 @@ export interface SessionInfo { currentOrgId: string; privileges: any; enableSpotterDataSourceDiscovery?: boolean; + // The access token surfaced by /prism/preauth/info (set server-side via + // sessionInfo.setAccessToken). Present for IAMv2/Okta sessions; used to mint + // org-scoped bearer tokens. May be undefined for sessions that don't expose one. + accessToken?: string; } export interface BaseMessage { diff --git a/static/oauth-callback.js b/static/oauth-callback.js index ee65d71..a4ef6b8 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,12 @@ } } - const data = await response.json(); + const raw = await response.json(); + // gettoken returns the access token at top-level `token` (with an optional + // `refreshToken`). Normalize to the { data: { token } } shape that + // /store-token expects (same shape the manual-paste path produces). + const accessToken = (raw && raw.data && raw.data.token) || raw.token; + const data = { data: { token: accessToken, refreshToken: raw.refreshToken } }; document.getElementById('status').textContent = 'Authentication successful. Securing your session...'; // Send the token to the server @@ -139,7 +144,7 @@ headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ + body: JSON.stringify({ token: data, oauthReqInfo: oauthReqInfo, instanceUrl: window.INSTANCE_URL