diff --git a/.gitignore b/.gitignore index e5c9c28..6dadd5c 100644 --- a/.gitignore +++ b/.gitignore @@ -137,3 +137,8 @@ dist *.egg-info __pycache__ .env*.local + +# Local chat history databases (mcp/python-react-agent-simple-ui) +*.db +*.db-wal +*.db-shm diff --git a/mcp/python-react-agent-simple-ui/README.md b/mcp/python-react-agent-simple-ui/README.md index 1e4adab..f88c1fb 100644 --- a/mcp/python-react-agent-simple-ui/README.md +++ b/mcp/python-react-agent-simple-ui/README.md @@ -1,6 +1,6 @@ # Python Agent with Simple React UI -A full-stack example that pairs a **Python (FastAPI) agent** with a **React chat UI**. Supports two backend implementations: +A full-stack example that pairs a **Python (FastAPI) agent** with a **React chat UI**, running against the ThoughtSpot MCP server's **Spotter 3** toolset. Two backends are included: -| Backend | File | AI Provider | MCP Integration | -|-----------------|----------------------------------------|-----------------------|------------------------------------------------| -| **v1 (OpenAI)** | `server/agent.py` | Azure OpenAI / OpenAI | Server-side (OpenAI manages MCP) | -| **v2 (Claude)** | `server/claude_agent_mcp_server_v2.py` | Anthropic Claude | Client-side (FastAPI connects to MCP directly) | -| **v2 (OpenAI)** | _(coming soon)_ | Azure OpenAI / OpenAI | Client-side (FastAPI connects to MCP directly) | +| Backend | File | What it gives you | +|-------------------------|--------------------------------------------------------------------|------------------------------------------------------------| +| **Spotter 3** | `server/claude_agent_with_spotter3_mcp_server.py` | The agent, with history held in memory for the process life | +| **Spotter 3 + history** | `server/claude_agent_with_spotter3_mcp_server_and_chat_history.py` | The same, plus SQLite history you can list, reopen and delete | - +Both use Anthropic Claude with a client-side MCP loop — the FastAPI process connects to the MCP server directly. The backend streams responses to the frontend using Server-Sent Events (SSE), giving users a real-time chat experience while the agent queries ThoughtSpot for data insights and displays ThoughtSpot charts in an embed. @@ -32,37 +33,112 @@ The backend streams responses to the frontend using Server-Sent Events (SSE), gi --- -## MCP Server v2: Claude + Client-side MCP (Recommended) +## Claude + the Spotter 3 MCP Server -`claude_agent_mcp_server_v2.py` uses Anthropic's Claude API with a **client-side agentic loop** — the FastAPI process connects directly to the ThoughtSpot MCP server using custom HTTP headers (`Authorization` + `x-ts-host`). This approach is required because Anthropic's server-side MCP integration does not support custom headers. +`claude_agent_with_spotter3_mcp_server.py` uses Anthropic's Claude API with a **client-side agentic loop** — the FastAPI process connects directly to the [ThoughtSpot MCP server](https://github.com/thoughtspot/mcp-server) using custom HTTP headers (`Authorization` + `x-ts-host`). This is required because Anthropic's server-side MCP connector cannot send custom headers. -### Architecture (v2) +### Architecture ``` ┌──────────────┐ SSE stream ┌──────────────────────┐ MCP (streamable-http) ┌─────────────┐ │ React Chat │ ◄────────────► │ FastAPI + Claude │ ◄──────────────────────────────► │ ThoughtSpot │ -│ (Vite) │ /api/chat │ / OpenAI │ Authorization + x-ts-host │ MCP Server │ +│ (Vite) │ /api/chat │ │ Authorization + x-ts-host │ MCP Server │ └──────────────┘ └──────────────────────┘ └─────────────┘ - :5173 :8000 agent.thoughtspot.app + :8000 :8001 agent.thoughtspot.app ``` **Request flow:** 1. User sends a message from the React UI -2. FastAPI opens a new MCP session to `agent.thoughtspot.app` with auth headers +2. FastAPI borrows the shared MCP session to `agent.thoughtspot.app` - opening it with auth headers and fetching the tool list only on the first message, or when the session has to be replaced (see [Shared MCP session](#shared-mcp-session)) 3. Claude receives the user message + ThoughtSpot tool definitions -4. Claude calls ThoughtSpot tools as needed; FastAPI executes each call via the MCP session -5. The agentic loop continues until Claude stops calling tools -6. Text deltas and status events are streamed to the UI over SSE in real time +4. Claude calls ThoughtSpot tools; FastAPI executes each call over the MCP session +5. For `get_session_updates`, FastAPI polls the Analytics Agent to completion itself (see below) +6. Text deltas, agent progress, and rendered answers stream to the UI over SSE in real time -### Prerequisites (v2) +### Integrating the ThoughtSpot MCP server into your own application -- Python 3.10+ +This example is meant to be copied from. The steps below are the path from "ThoughtSpot cluster" to "charts in my app's chat", each pointing at the code to lift. + +#### 1. Prepare ThoughtSpot and Anthropic + +- **A trusted-authentication secret key**, so your backend can mint tokens: *Develop > Customizations > Security Settings > Trusted authentication*. (A username + password also works for a demo, but MFA-enabled instances reject it.) +- **Your app's origin allowlisted for embedding.** The charts render in a ThoughtSpot iframe, and the cluster's CSP `frame-ancestors` must include the origin your UI is served from - otherwise the browser blocks the embed. This example's dev server is pinned to `http://localhost:8000` for exactly that reason (`client/vite.config.js`). See the [security settings docs](https://developers.thoughtspot.com/docs/security-settings) for CSP and CORS allowlists. +- **An Anthropic API key** for the agent. + +#### 2. Mint ThoughtSpot tokens on your server + +The secret key never reaches the browser. The backend calls `POST /api/rest/2.0/auth/token/full` and uses the result two ways: + +- `server_token()` - a cached token for the agent's own MCP and REST calls. A background task (`keep_server_token_fresh()`) mints it at startup and renews it about 12 minutes before it expires, so no request waits on a mint - which can take tens of seconds on a loaded cluster. +- `GET /api/ts-token` - a **fresh** token per call for the browser's Visual Embed SDK. The SDK requires a new token from every `getAuthToken` call; handing it the same one twice triggers its "Duplicate token" alert once that token stops verifying. + +Copy `mint_token()`, `server_token()` and the `/api/ts-token` endpoint. + +> **Production:** this example mints every token for one service user (`TS_EMBED_USERNAME`). In your app, mint for the **authenticated end user** instead, so ThoughtSpot's own permissions and row-level security apply to every question they ask. + +#### 3. Connect to the MCP server from your backend + +- Endpoint: `https://agent.thoughtspot.app/token/mcp?api-version=...`, with headers `Authorization: Bearer ` and `x-ts-host: `. +- Connect **client-side** (your process holds the MCP session), because Anthropic's server-side MCP connector cannot send the custom `x-ts-host` header. +- **Pin `api-version`** to a release date for anything you depend on - see [MCP endpoint and API version](#mcp-endpoint-and-api-version). +- **Share one session across requests.** Opening a session took 6-11 s on the clusters this was tested against. `connect_mcp()`, `McpPool` and `McpTurn` open it once, reconnect when it breaks or the token rotates, and retry a failed handshake once. + +#### 4. Hand the tools to your LLM + +- Convert `list_tools()` to your model's tool format (`build_tools()`), run the tool loop (`agent_loop()`), and return all parallel tool results in one message. +- **Poll `get_session_updates` in your code, not in the model.** `send_session_message` returns immediately and the Analytics Agent answers asynchronously; `autopoll_session_updates()` polls to completion and gives the model one consolidated result. +- **Strip `iframe_url` from what the model sees** (`strip_rendered_answers()`) - your UI renders the chart, so the model only needs to summarize it. + +#### 5. Stream progress and answers to your UI + +Stream text, progress and answers as they arrive - an Analytics Agent answer typically takes tens of seconds. This example uses Server-Sent Events; the event contract is in [SSE events](#sse-events). + +#### 6. Render the charts with the Visual Embed SDK + +```ts +init({ + thoughtSpotHost: import.meta.env.VITE_TS_HOST, + authType: AuthType.TrustedAuthTokenCookieless, + getAuthToken: async () => { + const response = await fetch("/api/ts-token", { cache: "no-store" }); + if (!response.ok) throw new Error(`Token request failed: ${response.status}`); + return (await response.json()).token; + }, + autoLogin: true, // fetch a replacement before the current token expires +}); + +// Upgrades every `; +}; + +/** + * One MCP answer's iframe, written into the DOM exactly once. + * + * startAutoMCPFrameRenderer swaps the iframe for the real embed via replaceWith(), so + * React must not own that node - it owns only this wrapper, filled with + * `dangerouslySetInnerHTML`. React re-applies that prop whenever it receives a new + * object, and a fresh `{ __html }` literal is a new object on every render of the message + * list - each keystroke, each streamed delta. Every re-application throws the embed away + * and the renderer resolves the answer again from scratch: two ThoughtSpot calls per + * chart, and a stored answer's take tens of seconds. + * + * So the prop object is created once per mount and kept in state (which also survives + * dev Fast Refresh). Callers key this by `html`, so a different answer gets a fresh + * wrapper. The markup goes in before the wrapper is attached, so the renderer sees the + * iframe once - writing it from an effect instead would show it twice, via the attached + * container and via the iframe itself, and resolve it twice. + */ +function AnswerFrame({ html }: { html: string }) { + const [markup] = useState(() => ({ __html: html })); + return
; +} + +function App() { + const [messages, setMessages] = useState([]); + const [input, setInput] = useState(""); + const [isLoading, setIsLoading] = useState(false); + const [status, setStatus] = useState(""); + const [responseId, setResponseId] = useState(null); + // ThoughtSpot analytical session for the open conversation. Only a reopened + // chat needs it - live answers arrive with a URL already attached. + const [sessionId, setSessionId] = useState(null); + // Chat history. `historyAvailable` stays false against a backend without the + // /api/conversations endpoints, and the sidebar simply isn't rendered. + const [conversations, setConversations] = useState([]); + const [historyAvailable, setHistoryAvailable] = useState(false); + // The stored chat being fetched. Opening one waits on ThoughtSpot (the server + // reconciles its answers against the live conversation), so it can take seconds. + const [openingId, setOpeningId] = useState(null); + // Bumped on every open / new chat, so only the latest request's response lands. + const openRequestRef = useRef(0); + const messagesEndRef = useRef(null); + const textareaRef = useRef(null); + + const refreshConversations = useCallback(async () => { + try { + const response = await fetch(HISTORY_URL); + if (!response.ok) throw new Error(String(response.status)); + const data = (await response.json()) as { + conversations?: ConversationSummary[]; + }; + setConversations(data.conversations || []); + setHistoryAvailable(true); + } catch { + setHistoryAvailable(false); + } + }, []); + + useEffect(() => { + refreshConversations(); + }, [refreshConversations]); + + // Reopen a stored conversation. Stored turns carry their answers, so the charts + // come back as embeds rather than as text - resolved live, since the stored + // answers deliberately carry no URL. See answerSrc. + const openConversation = useCallback( + async (id: string) => { + if (isLoading) return; + const request = ++openRequestRef.current; + setOpeningId(id); + setStatus(""); + try { + const response = await fetch(`${HISTORY_URL}/${id}`); + if (!response.ok) throw new Error(String(response.status)); + const data = (await response.json()) as ConversationDetail; + // Another chat was opened, or a new one started, while this was loading. + if (request !== openRequestRef.current) return; + setMessages( + (data.turns || []).map((turn) => ({ + role: turn.role, + content: turn.content, + answers: turn.answers || [], + work: turn.work || [], + workAnswerIds: turn.work_answer_ids || [], + })), + ); + setResponseId(data.id); + setSessionId(data.analytical_session_id || null); + setStatus(""); + } catch (error) { + if (request !== openRequestRef.current) return; + setStatus(`Could not open that conversation: ${errorMessage(error)}`); + } finally { + if (request === openRequestRef.current) setOpeningId(null); + } + }, + [isLoading], + ); + + const deleteConversation = useCallback( + async (id: string, event: MouseEvent) => { + event.stopPropagation(); + if (id === openingId) { + // Deleting the chat that is still loading: drop the load too. + openRequestRef.current++; + setOpeningId(null); + } + await fetch(`${HISTORY_URL}/${id}`, { method: "DELETE" }); + if (id === responseId) { + setMessages([]); + setResponseId(null); + } + refreshConversations(); + }, + [responseId, openingId, refreshConversations], + ); + + const scrollToBottom = useCallback(() => { + messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); + }, []); + + useEffect(() => { + scrollToBottom(); + }, [messages, status, scrollToBottom]); + + useEffect(() => { + if (textareaRef.current) { + textareaRef.current.style.height = "auto"; + textareaRef.current.style.height = + Math.min(textareaRef.current.scrollHeight, 150) + "px"; + } + }, [input]); + + const sendMessage = async () => { + // While a stored chat loads, a send would go to the previous conversation and + // its reply would land in the one being opened. + if (!input.trim() || isLoading || openingId) return; + + const userMessage = input.trim(); + setInput(""); + setIsLoading(true); + setStatus("Thinking..."); + + setMessages((prev) => [ + ...prev, + { role: "user", content: userMessage }, + { role: "assistant", content: "", answers: [], work: [] }, + ]); + + try { + const response = await fetch(API_URL, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + message: userMessage, + response_id: responseId, + }), + }); + + if (!response.ok) { + throw new Error(`Server error: ${response.status}`); + } + + if (!response.body) throw new Error("Server sent an empty response"); + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split("\n"); + buffer = lines.pop() || ""; + + for (const line of lines) { + if (!line.startsWith("data: ")) continue; + + try { + const data = JSON.parse(line.slice(6)) as StreamEvent; + + switch (data.type) { + case "delta": + setStatus(""); + setMessages((prev) => { + const updated = [...prev]; + const last = updated[updated.length - 1]; + updated[updated.length - 1] = { + ...last, + content: last.content + data.text, + }; + return updated; + }); + break; + case "status": + setStatus(data.message); + break; + // An answer the Analytics Agent produced. The server streams these as + // soon as they arrive, so charts appear while the agent is still working. + case "answer": { + // Raw MCP updates have no `iframe_url`, so an answer is renderable + // as long as a src can be built from what it does carry. + const src = answerSrc(data, null); + if (!src) break; + setMessages((prev) => { + const updated = [...prev]; + const last = updated[updated.length - 1]; + const answers = last.answers || []; + if (answers.some((a) => answerSrc(a, null) === src)) return prev; + updated[updated.length - 1] = { + ...last, + answers: [...answers, data], + }; + return updated; + }); + break; + } + case "work": + setMessages((prev) => { + const updated = [...prev]; + const last = updated[updated.length - 1]; + updated[updated.length - 1] = { + ...last, + work: appendWork(last.work || [], data), + }; + return updated; + }); + break; + case "done": + setResponseId(data.response_id); + setStatus(""); + // The server has written the turn by the time it sends `done`. + refreshConversations(); + break; + case "error": + setMessages((prev) => { + const updated = [...prev]; + const last = updated[updated.length - 1]; + updated[updated.length - 1] = { + role: "assistant", + content: `Error: ${data.message}`, + answers: last.answers || [], + work: last.work || [], + isError: true, + }; + return updated; + }); + break; + } + } catch { + /* skip malformed lines */ + } + } + } + } catch (error) { + setMessages((prev) => { + const updated = [...prev]; + if (updated.length > 0) { + updated[updated.length - 1] = { + role: "assistant", + content: `Failed to connect to server: ${errorMessage(error)}`, + isError: true, + }; + } + return updated; + }); + } finally { + setIsLoading(false); + setStatus(""); + } + }; + + const handleKeyDown = (e: KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + sendMessage(); + } + }; + + const startNewChat = () => { + // Drop any chat still loading, so it doesn't replace the new one when it lands. + openRequestRef.current++; + setOpeningId(null); + setMessages([]); + setResponseId(null); + setSessionId(null); + setStatus(""); + setInput(""); + }; + + return ( +
+ {historyAvailable && ( + + )} + +
+
+
+

ThoughtSpot Agent

+ {messages.length > 0 && ( + + )} +
+
+ +
+ {openingId ? ( +
+
+ ) : messages.length === 0 ? ( +
+
TS
+

Ask anything about your data

+

+ Powered by ThoughtSpot and Claude. Ask questions and get + insights from your connected data sources. +

+
+ ) : ( +
+ {messages.map((msg, i) => ( +
+
+ {msg.role === "user" ? "U" : "TS"} +
+
+ {msg.role !== "assistant" ? ( +

{msg.content}

+ ) : ( + <> + {!!msg.work?.length && ( + + )} + {(msg.answers || []).map((answer) => ( +
+ {answer.title && ( +
{answer.title}
+ )} + {(() => { + const html = answerHtml(answer, sessionId); + return ; + })()} +
+ ))} + {msg.content ? ( + // rehypeRaw so an " # important! - "Do not ask to create charts, as thoughtspot will already create interactive charts for you." - "Respond in an engaging markdown format, with html tags when needed." - "Keep the response short and to the point." - # "Use this datasource: cd252e5c-b552-49a8-821d-3eadaa049cca to answer all data questions." -) - -# ThoughtSpot MCP Tools (v2): -# To restrict which tools are accessible to the agent, set ALLOWED_TOOLS to a list of tool names. -# Set to None to allow all tools. -# -# The v2 MCP server uses an analytical session workflow: -# 1. check_connectivity - Test connectivity and authentication. No inputs. -# 2. create_analysis_session - Start a session. Optional: data_source_id. -# Returns: analytical_session_id. -# 3. send_session_message - Send a natural-language question to the session. -# Inputs: analytical_session_id, message, additional_context (optional). -# 4. get_session_updates - Poll for incremental updates. Inputs: analytical_session_id. -# Returns: session_updates (list), is_done (bool). -# Poll until is_done=True. Each update has type: text | text_chunk | answer. -# Answer updates include: answer_id, answer_title, answer_query, iframe_url. -# 5. create_dashboard - Create a dashboard from answer IDs. -# Inputs: title, answers (list of answer_ids), note_tile. -# Returns: link. -ALLOWED_TOOLS = None -# ALLOWED_TOOLS = ["check_connectivity", "create_analysis_session", "send_session_message", "get_session_updates", "create_dashboard"] - -# In-memory conversation store: conv_id -> full message history (including tool interactions) -conversations: dict[str, list] = {} - -# 1:1 mapping: conv_id -> analytical_session_id returned by create_analysis_session tool. -# Passed to Claude via system prompt so follow-up send_session_message / get_session_updates -# calls use the same ThoughtSpot analytical session. -analytical_sessions: dict[str, str] = {} - - -class ChatRequest(BaseModel): - message: str - response_id: str | None = None - - -def format_sse(data: dict) -> str: - return f"data: {json.dumps(data)}\n\n" - - -async def agent_loop(messages: list, queue: asyncio.Queue, conv_id: str) -> None: - """ - Client-side agentic loop. Connects to the ThoughtSpot MCP server directly - (with Authorization + x-ts-host headers), fetches tool definitions, then - runs the Claude tool-use loop until the model stops calling tools. - Puts SSE event dicts into queue for streaming to the frontend. - """ - try: - headers = dict(MCP_HEADERS) - - print(f"[MCP] Connecting to {MCP_URL}") - async with streamablehttp_client(MCP_URL, headers=headers) as (read, write, _): - async with ClientSession(read, write) as session: - print("[MCP] Initializing session...") - await session.initialize() - print("[MCP] Session initialized. Fetching tools...") - - # Fetch tool definitions from ThoughtSpot MCP server - tools_result = await session.list_tools() - print(f"[MCP] Got {len(tools_result.tools)} tools") - available_tools = tools_result.tools - # Optionally filter tools based on ALLOWED_TOOLS - if ALLOWED_TOOLS is not None: - available_tools = [t for t in available_tools if t.name in ALLOWED_TOOLS] - - # Convert MCP tool definitions to Anthropic format - anthropic_tools = [ - { - "name": t.name, - "description": t.description or "", - "input_schema": t.inputSchema, - } - for t in available_tools - ] - - current_messages = messages[:] - final_text_parts: list[str] = [] - - # Build system prompt, injecting analytical_session_id for follow-up turns - system = SYSTEM_PROMPT - existing_session_id = analytical_sessions.get(conv_id) - if existing_session_id: - system += ( - f"\n\nActive ThoughtSpot analytical session ID: {existing_session_id}. " - "Use this ID when calling send_session_message or get_session_updates " - "so follow-up questions continue in the same session." - ) - - while True: - async with claude_client.messages.stream( - model="claude-opus-4-6", - max_tokens=16000, - system=system, - messages=current_messages, - tools=anthropic_tools, - ) as stream: - async for event in stream: - t = getattr(event, "type", None) - if t == "content_block_start": - if getattr(event.content_block, "type", None) == "tool_use": - await queue.put({"type": "status", "message": "Querying ThoughtSpot..."}) - elif t == "content_block_delta": - delta = event.delta - if getattr(delta, "type", None) == "text_delta": - await queue.put({"type": "delta", "text": delta.text}) - final_text_parts.append(delta.text) - - final_message = await stream.get_final_message() - - if final_message.stop_reason != "tool_use": - break - - # Execute each tool call via MCP client (headers are set on the session) - tool_results = [] - for block in final_message.content: - if getattr(block, "type", None) == "tool_use": - try: - mcp_result = await session.call_tool(block.name, block.input) - print(f"[MCP] Tool {block.name} and input {block.input} returned: {mcp_result}") - result_text = " ".join( - getattr(c, "text", str(c)) for c in mcp_result.content - ) if mcp_result.content else "" - is_error = getattr(mcp_result, "isError", False) - - # Store analytical_session_id (1:1 with conv_id) so follow-up - # requests can reference the same ThoughtSpot session. - if block.name == "create_analysis_session" and not analytical_sessions.get(conv_id): - try: - sid = json.loads(result_text).get("analytical_session_id") - if sid: - analytical_sessions[conv_id] = sid - print(f"[MCP] Stored analytical_session_id for conv {conv_id}: {sid}") - except Exception: - pass - - except McpError as e: - print(f"[MCP] Tool {block.name} failed: {e}") - result_text = f"Tool call failed: {e}" - is_error = True - tool_results.append({ - "type": "tool_result", - "tool_use_id": block.id, - "content": result_text, - "is_error": is_error, - }) - - # Append assistant turn + tool results and continue the loop - current_messages = current_messages + [ - {"role": "assistant", "content": final_message.content}, - {"role": "user", "content": tool_results}, - ] - final_text_parts = [] # reset; next iteration may stream more text - - # Persist full conversation history (including tool interactions) so - # follow-up turns have complete context (e.g. analytical_session_id in prior results). - conversations[conv_id] = current_messages + [ - {"role": "assistant", "content": final_message.content} - ] - await queue.put({"type": "done", "response_id": conv_id}) - - except BaseException as e: - traceback.print_exc() - # Recursively unwrap ExceptionGroup to get the root cause - err = e - while hasattr(err, "exceptions") and getattr(err, "exceptions", None): - err = err.exceptions[0] - await queue.put({"type": "error", "message": f"{type(err).__name__}: {err}"}) - - -@app.post("/api/chat") -async def chat(request: ChatRequest): - conv_id = request.response_id or str(uuid.uuid4()) - print(f"[Chat] Received message for conv_id {conv_id}: {request.response_id}") - history = conversations.get(conv_id, []) - messages = history + [{"role": "user", "content": request.message}] - - queue: asyncio.Queue = asyncio.Queue() - asyncio.create_task(agent_loop(messages, queue, conv_id)) - - async def event_stream() -> AsyncGenerator[str, None]: - while True: - item = await queue.get() - yield format_sse(item) - if item.get("type") in ("done", "error"): - break - - return StreamingResponse(event_stream(), media_type="text/event-stream") - - -@app.get("/api/health") -async def health(): - return {"status": "ok"} diff --git a/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server.py b/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server.py new file mode 100644 index 0000000..cdecaeb --- /dev/null +++ b/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server.py @@ -0,0 +1,1104 @@ +""" +Python agent: Anthropic Claude + the ThoughtSpot MCP server (Spotter 3 toolset). + +FastAPI streams chat responses to the React frontend over Server-Sent Events. + +Why client-side MCP: the ThoughtSpot MCP server's static-token endpoint needs custom +HTTP headers (Authorization + x-ts-host). Anthropic's server-side MCP connector cannot +send those, so this process connects to the MCP server itself and executes tool calls. + +Two things this server does that a plain pass-through loop does not: + +1. `get_session_updates` polling happens here, not in the model. The ThoughtSpot + Analytics Agent answers asynchronously, so one `get_session_updates` call usually + returns `is_done: false` and an empty list. Letting the model poll costs a full + model round-trip per poll. Instead we poll until `is_done: true` and hand the model + one consolidated result, streaming the Agent's progress to the UI as it arrives. + +2. Answers are rendered by the client, not by the model. Each `answer` update becomes + an iframe marked `tsmcp=true`, which the React client mounts and the Visual Embed + SDK's `startAutoMCPFrameRenderer` upgrades into a real ThoughtSpot embed. We keep + the URL out of what the model sees - it is long, and the model does not need to + echo markup for a chart the UI has already drawn. + +MCP server: https://github.com/thoughtspot/mcp-server +""" + +import asyncio +import json +import os +import time +import traceback +import uuid +from collections.abc import AsyncGenerator, AsyncIterator, Callable +from contextlib import AsyncExitStack, asynccontextmanager +from pathlib import Path +from typing import Any + +import anthropic +import httpx2 +from dotenv import load_dotenv +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import StreamingResponse +from mcp import Client, MCPError +from mcp.client.streamable_http import streamable_http_client +from pydantic import BaseModel + +load_dotenv(dotenv_path=Path(__file__).resolve().parent.parent / ".env") + +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + refresher = asyncio.create_task(keep_server_token_fresh()) if CAN_MINT_TOKENS else None + yield + if refresher: + refresher.cancel() + # Close the shared MCP session so the server side is not left holding it. + await mcp_pool.close() + + +app = FastAPI(title="ThoughtSpot Agent", lifespan=lifespan) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# ── Anthropic ─────────────────────────────────────────────────────────────────── +claude_client = anthropic.AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + +# Haiku 4.5: the agent only orchestrates ThoughtSpot's tools - pick a tool, pass the +# question on, summarise the answer - and the Analytics Agent does the analysis. Measured +# on one question, Haiku's four model calls took ~4s in total, against ~17s on Sonnet 5 +# and ~8s on Opus 5. Set ANTHROPIC_MODEL=claude-sonnet-5 for stronger reasoning on +# multi-step follow-ups, or claude-opus-5 for the most capable. +MODEL = os.getenv("ANTHROPIC_MODEL", "claude-haiku-4-5") +MAX_TOKENS = 16000 + +# Safety classifiers can decline a request (HTTP 200, stop_reason="refusal"). The +# server-side fallback beta re-routes those to a fallback model automatically - but +# only models with published fallback targets accept it. Sonnet 5 and Haiku 4.5 have +# none (`allowed_fallback_models` is empty on /v1/models), so for them a refusal is +# simply reported to the user. +REFUSAL_FALLBACK_BETA = "server-side-fallback-2026-07-01" +FALLBACK_CAPABLE_MODELS = ("claude-opus-5", "claude-fable-5-1") + + +def model_request_options(model: str) -> dict: + """The thinking and fallback options this model accepts.""" + options: dict[str, Any] = {} + # Haiku 4.5 predates adaptive thinking (it would need a fixed `budget_tokens`). + # Choosing tools does not need extended reasoning, so run it without thinking - + # that is also the fastest configuration. + if not model.startswith("claude-haiku"): + options["thinking"] = {"type": "adaptive", "display": "summarized"} + if model in FALLBACK_CAPABLE_MODELS: + options["betas"] = [REFUSAL_FALLBACK_BETA] + options["fallbacks"] = "default" + return options + +# ── ThoughtSpot MCP server ────────────────────────────────────────────────────── +TS_HOST = os.getenv("VITE_TS_HOST") or os.getenv("TS_HOST") +TS_AUTH_TOKEN = os.getenv("VITE_TS_AUTH_TOKEN") or os.getenv("TS_AUTH_TOKEN") + +if not TS_HOST: + raise RuntimeError("TS_HOST (or VITE_TS_HOST) must be set in .env") + +# The `/token/*` endpoint family is the static-bearer-token transport. +# (`/bearer/*` is the legacy path and is frozen on the older v1 toolset.) +# +# `latest` tracks the newest toolset, so a ThoughtSpot release can change the tools +# and the update shape under this app. Set TS_MCP_API_VERSION to a release date +# (e.g. 2026-05-01) to pin it instead. +MCP_API_VERSION = os.getenv("TS_MCP_API_VERSION", "latest") +# +# Adding `&enable-raw-session-updates=true` makes the server stream the Agent's own +# updates instead of its digested ones. This app reads either - see "Reading session +# updates" below - so you can turn it on through TS_MCP_URL without touching the code. +MCP_URL = os.getenv( + "TS_MCP_URL", + f"https://agent.thoughtspot.app/token/mcp?api-version={MCP_API_VERSION}", +) + +# Headers for the MCP session are built per request by `mcp_headers()` below, so a +# long-running process picks up a renewed token instead of holding an expired one. + +# ── ThoughtSpot tokens ────────────────────────────────────────────────────────── +# Two consumers, two policies: +# +# The browser. The Visual Embed SDK requires a FRESH token from every `getAuthToken` +# call. Handing it the same static token twice trips its duplicate-token check the +# moment the token stops verifying, and the callback can never recover because it +# returns the same string. So `/api/ts-token` mints one per request. +# +# This server, for its own MCP and REST calls. A static TS_AUTH_TOKEN eventually +# expires, and when it does every tool call fails with "Failed to validate +# connection" while the tool *list* still succeeds - so the app looks healthy right +# up until someone asks a question. Instead we mint one and cache it until shortly +# before it expires. TS_AUTH_TOKEN stays as the fallback for when no minting +# credentials are configured. +# +# Minting needs a username plus either the cluster secret key (Develop > Customizations +# > Security Settings > Trusted authentication) or that user's password. `secret_key` +# takes precedence when both are set. +TS_EMBED_USERNAME = os.getenv("TS_EMBED_USERNAME") +TS_SECRET_KEY = os.getenv("TS_SECRET_KEY") +TS_EMBED_PASSWORD = os.getenv("TS_EMBED_PASSWORD") +TS_TOKEN_VALIDITY_SEC = int(os.getenv("TS_TOKEN_VALIDITY_SEC", "1800")) + +CAN_MINT_TOKENS = bool(TS_EMBED_USERNAME and (TS_SECRET_KEY or TS_EMBED_PASSWORD)) + +if not CAN_MINT_TOKENS and not TS_AUTH_TOKEN: + raise RuntimeError( + "Set TS_EMBED_USERNAME plus TS_SECRET_KEY (or TS_EMBED_PASSWORD) so the server can " + "mint ThoughtSpot tokens, or a static TS_AUTH_TOKEN (or VITE_TS_AUTH_TOKEN) in .env" + ) + +if not CAN_MINT_TOKENS: + print( + "[Auth] TS_EMBED_USERNAME + TS_SECRET_KEY (or TS_EMBED_PASSWORD) are not set - " + "/api/ts-token will serve the static TS_AUTH_TOKEN. Fine for a local demo, but " + "the SDK cannot recover once that token expires." + ) + + +# Renew this long before the token actually expires, so a call already in flight does +# not race the expiry. +TOKEN_REFRESH_MARGIN_SEC = 120 +# The server's own token. Longer-lived than the browser's: it is never handed out, and +# a longer life means fewer mints on a slow cluster. +SERVER_TOKEN_VALIDITY_SEC = int(os.getenv("TS_SERVER_TOKEN_VALIDITY_SEC", "3600")) + +_server_token: str | None = None +# Wall-clock (time.time()) instant to renew at. Not time.monotonic(): on macOS that clock +# stops while the machine sleeps, so after a laptop sleep the cache would keep serving a +# token the cluster has already expired. +_server_token_expiry = 0.0 +# When the background task renews the token - well before `_server_token_expiry`, so no +# request has to wait on a mint. Minting can take 40s on a loaded cluster, and a request +# that finds the token expired pays all of it. +_server_token_refresh_at = 0.0 +# Serialises minting: an agent turn fires several tool calls at once, and without this +# they would each mint a token on a cold cache. +_server_token_lock = asyncio.Lock() + +# Renew this long before the token expires. Capped at half the token's life, so a short +# cluster-imposed lifetime still gets renewed at a sane interval. +TOKEN_BACKGROUND_LEAD_SEC = 720 +# After a failed background mint, try again this soon. The current token keeps serving. +TOKEN_BACKGROUND_RETRY_SEC = 30.0 +# The background loop re-reads the wall clock at least this often, so it notices when a +# laptop sleep has carried it past the renewal time (asyncio.sleep pauses with the host). +TOKEN_BACKGROUND_TICK_SEC = 60.0 + + +class TokenMintError(Exception): + """The cluster would not issue a token.""" + + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message) + self.status_code = status_code + + +async def mint_token(validity_sec: int) -> dict: + """Ask ThoughtSpot for a bearer token. Returns the parsed response body.""" + payload: dict[str, Any] = { + "username": TS_EMBED_USERNAME, + "validity_time_in_sec": validity_sec, + } + # secret_key wins over password when both are present, per the REST API contract. + if TS_SECRET_KEY: + payload["secret_key"] = TS_SECRET_KEY + else: + payload["password"] = TS_EMBED_PASSWORD + + # This endpoint is occasionally slow to respond - 15s is normal on a loaded + # cluster - and a bare timeout would surface as an unhandled 500 with no + # explanation. + try: + async with httpx2.AsyncClient(timeout=httpx2.Timeout(10.0, read=90.0)) as client: + response = await client.post( + f"{TS_HOST.rstrip('/')}/api/rest/2.0/auth/token/full", json=payload + ) + except httpx2.HTTPError as exc: + raise TokenMintError(f"{type(exc).__name__}", 504) from exc + + if response.status_code != 200: + # Surface the cluster's own reason (bad secret key, trusted auth disabled, ...) + # rather than a blank 500 the browser cannot explain. + print(f"[Auth] Token mint failed {response.status_code}: {response.text[:300]}") + raise TokenMintError(f"HTTP {response.status_code}", 502) + + return response.json() + + +def server_token_valid() -> bool: + return bool(_server_token) and time.time() < _server_token_expiry + + +async def mint_server_token() -> str: + """Mint a server token and cache it. The caller holds `_server_token_lock`.""" + global _server_token, _server_token_expiry, _server_token_refresh_at + + body = await mint_token(SERVER_TOKEN_VALIDITY_SEC) + created = body.get("creation_time_in_millis") + expires = body.get("expiration_time_in_millis") + lifetime = ( + (expires - created) / 1000 + if isinstance(created, (int, float)) and isinstance(expires, (int, float)) + else SERVER_TOKEN_VALIDITY_SEC + ) + now = time.time() + _server_token = body["token"] + _server_token_expiry = now + max(60.0, lifetime - TOKEN_REFRESH_MARGIN_SEC) + _server_token_refresh_at = now + max(lifetime / 2, lifetime - TOKEN_BACKGROUND_LEAD_SEC) + print(f"[Auth] Minted a server token, good for {lifetime:.0f}s", flush=True) + return _server_token + + +async def server_token() -> str: + """A currently-valid bearer token for this server's own ThoughtSpot calls.""" + if not CAN_MINT_TOKENS: + return TS_AUTH_TOKEN + + # Fast path, outside the lock: a valid token never waits on a mint in progress - + # the background renewal holds the lock for as long as the cluster takes. + if server_token_valid(): + return _server_token + + async with _server_token_lock: + if server_token_valid(): + return _server_token + try: + return await mint_server_token() + except TokenMintError as exc: + # Keep serving the token we have: an expiring token still beats no token, + # and the next call retries the mint. + fallback = _server_token or TS_AUTH_TOKEN + if not fallback: + raise + print(f"[Auth] Could not mint a server token ({exc}); using the previous one.") + return fallback + + +async def keep_server_token_fresh() -> None: + """Background task: mint at startup, then renew ahead of expiry. + + Renewing on demand still works - this only moves the wait off the request path. + """ + while True: + due = _server_token_refresh_at if server_token_valid() else 0.0 + remaining = due - time.time() + if remaining > 0: + await asyncio.sleep(min(TOKEN_BACKGROUND_TICK_SEC, remaining)) + continue + try: + async with _server_token_lock: + # A request may have minted while this task waited for the lock. + if server_token_valid() and time.time() < _server_token_refresh_at: + continue + await mint_server_token() + except TokenMintError as exc: + print(f"[Auth] Background token renewal failed ({exc}); retrying soon.") + await asyncio.sleep(TOKEN_BACKGROUND_RETRY_SEC) + except Exception: # noqa: BLE001 - the loop must outlive any one failure + traceback.print_exc() + await asyncio.sleep(TOKEN_BACKGROUND_RETRY_SEC) + + +def invalidate_server_token() -> None: + """Force the next server_token() call to mint, keeping the old token as a fallback.""" + global _server_token_expiry, _server_token_refresh_at + _server_token_expiry = 0.0 + _server_token_refresh_at = 0.0 + + +async def mcp_headers() -> dict: + """Auth headers for an MCP session, with a token that is current right now.""" + return {"Authorization": f"Bearer {await server_token()}", "x-ts-host": TS_HOST} + +# Long read timeout: MCP replies stream over SSE and the Analytics Agent is slow. +MCP_TIMEOUT = httpx2.Timeout(30.0, read=300.0) + +# The MCP server's `initialize` intermittently fails with a bare HTTP 500 (surfaced as +# "MCPError: Server returned an error response") that a second attempt gets past. So a +# failed handshake is retried once. The retry also mints a fresh token - cheap, and it +# rules the token out. Only the handshake is retried: past it, the turn has streamed +# output to the browser and cannot be replayed. +MCP_CONNECT_ATTEMPTS = 2 +MCP_CONNECT_RETRY_DELAY_SEC = 2.0 + + +def root_cause(exc: BaseException) -> BaseException: + """anyio wraps transport failures in ExceptionGroups; unwrap to the first leaf.""" + while getattr(exc, "exceptions", None): + exc = exc.exceptions[0] + return exc + + +@asynccontextmanager +async def connect_mcp( + send_status: Callable[[dict], None] | None = None, +) -> AsyncIterator[tuple[Client, Any, str]]: + """An initialised MCP session, its tool listing and the bearer token it runs on. + + Retries a failed handshake once. + """ + for attempt in range(1, MCP_CONNECT_ATTEMPTS + 1): + stack = AsyncExitStack() + try: + headers = await mcp_headers() + token = headers["Authorization"].removeprefix("Bearer ") + http_client = await stack.enter_async_context( + httpx2.AsyncClient(headers=headers, timeout=MCP_TIMEOUT, follow_redirects=True) + ) + mcp = await stack.enter_async_context( + Client(streamable_http_client(MCP_URL, http_client=http_client)) + ) + listing = await mcp.list_tools() + except BaseException as exc: + await stack.aclose() + if not isinstance(exc, Exception) or attempt == MCP_CONNECT_ATTEMPTS: + raise + err = root_cause(exc) + print( + f"[MCP] Handshake attempt {attempt} failed ({type(err).__name__}: {err}); " + "retrying with a fresh token" + ) + if send_status: + send_status({"type": "status", "message": "Reconnecting to ThoughtSpot..."}) + invalidate_server_token() + await asyncio.sleep(MCP_CONNECT_RETRY_DELAY_SEC) + continue + + async with stack: + yield mcp, listing, token + return + + +# ── Shared MCP session ────────────────────────────────────────────────────────── +# +# Opening an MCP session costs ~10s on a loaded cluster (discover probe, initialize, +# tools/list - each several seconds), which used to be paid on every chat message. One +# session is now shared by every turn and reopened only when it has to be: +# +# - the server token rotated (the session's HTTP client carries the old one), or +# - the session broke: a transport error, or the server dropped it (-32600 +# "Session terminated", e.g. after an idle timeout). +# +# The session is opened and closed by its own background task, never by a request. +# anyio requires a session to be closed by the task that opened it, and closing it in +# a request task is what used to get cancelled half-way when the browser stream ended. + +# What the MCP client raises once the server no longer knows the session: -32600 with +# the message "Session terminated". Only that pairing means the call never ran and is +# safe to repeat. -32600 alone is not enough - it is JSON-RPC's generic "Invalid +# Request", which the client also uses when a call may already have been processed. +MCP_SESSION_TERMINATED = (-32600, "Session terminated") + + +def is_session_terminated(exc: MCPError) -> bool: + return (exc.code, exc.message) == MCP_SESSION_TERMINATED + + +class McpConnection: + """One open MCP session, held open by its owner task until retired. + + `leases` counts the turns using it. A retired connection stays open until the last + of them is done, so rotating the token never pulls a session out from under a turn. + """ + + def __init__( + self, mcp: Client, listing: Any, token: str, stop: asyncio.Event, task: asyncio.Task + ) -> None: + self.mcp = mcp + self.listing = listing + self.token = token + self.task = task + self._stop = stop + self.leases = 0 + self.retired = False + + @property + def usable(self) -> bool: + return not self.retired and not self.task.done() + + # Both of these are synchronous on purpose: a turn gives up its lease without + # awaiting anything, so a cancelled turn cannot be interrupted mid-teardown. The + # owner task does the actual closing. + def retire(self) -> None: + self.retired = True + if self.leases == 0: + self._stop.set() + + def release(self) -> None: + self.leases -= 1 + if self.retired and self.leases == 0: + self._stop.set() + + +class McpPool: + """Hands out the shared MCP session, opening a new one when the current can't serve.""" + + def __init__(self) -> None: + # Held across get-or-create, so turns arriving together share one handshake + # instead of each opening a session. + self._lock = asyncio.Lock() + self._conn: McpConnection | None = None + + async def acquire(self, send_status: Callable[[dict], None] | None = None) -> McpConnection: + async with self._lock: + token = await server_token() + conn = self._conn + if conn and conn.usable and conn.token == token: + conn.leases += 1 + return conn + if conn: + reason = "token rotated" if conn.usable else "session closed" + print(f"[MCP] Replacing shared session ({reason})") + conn.retire() + self._conn = None + conn = await self._open(send_status) + conn.leases += 1 + self._conn = conn + return conn + + async def _open(self, send_status: Callable[[dict], None] | None) -> McpConnection: + ready: asyncio.Future = asyncio.get_running_loop().create_future() + stop = asyncio.Event() + + async def own() -> None: + try: + async with connect_mcp(send_status) as opened: + if ready.done(): + return # the asker left mid-handshake; just let the session close + ready.set_result(opened) + await stop.wait() + except Exception as exc: # noqa: BLE001 - reported to the waiter, or logged + if not ready.done(): + ready.set_exception(exc) + else: + err = root_cause(exc) + print(f"[MCP] Shared session ended: {type(err).__name__}: {err}") + except BaseException: + # Cancellation, possibly grouped with other errors by anyio. Resolve the + # waiter so a turn never waits forever on a handshake that is gone. + if not ready.done(): + ready.cancel() + raise + + task = asyncio.create_task(own()) + try: + mcp, listing, token = await ready + except asyncio.CancelledError: + # Close the session once it opens rather than leave it with no owner. + stop.set() + if not asyncio.current_task().cancelling(): + # This turn was not cancelled - the handshake was (e.g. at shutdown). + # Report it as a failure, not as a hang-up nobody is told about. + raise RuntimeError("The MCP session closed while it was opening") from None + raise + return McpConnection(mcp, listing, token, stop, task) + + async def close(self) -> None: + """Retire the current session and wait for it to close (app shutdown).""" + async with self._lock: + conn, self._conn = self._conn, None + if conn: + conn.leases = 0 + conn.retire() + await asyncio.wait({conn.task}, timeout=10) + + +mcp_pool = McpPool() + + +class McpTurn: + """One chat turn's handle on the shared MCP session. + + Exposes `call_tool` with Client's signature, so the tool code does not know the + session is shared. A call that finds the session dropped is repeated once on a + fresh one; any other connection failure retires the session for the next turn + and surfaces the error, since the call may already have run. + """ + + def __init__(self, send_status: Callable[[dict], None] | None = None) -> None: + self.send_status = send_status + self.conn: McpConnection | None = None + # Every connection this turn has leased. A replaced one is released only when + # the turn ends, so a sibling call still in flight on it is not cut off. + self._held: list[McpConnection] = [] + # Tool calls run concurrently; only one of them should swap the session. + self._swap_lock = asyncio.Lock() + + async def __aenter__(self) -> "McpTurn": + self.conn = await mcp_pool.acquire(self.send_status) + self._held.append(self.conn) + return self + + async def __aexit__(self, *exc_info: Any) -> None: + for conn in self._held: + conn.release() + self._held.clear() + self.conn = None + + @property + def listing(self) -> Any: + return self.conn.listing + + async def call_tool(self, name: str, arguments: dict) -> Any: + for attempt in (1, 2): + conn = self.conn + try: + return await conn.mcp.call_tool(name, arguments) + except MCPError as exc: + if attempt == 2 or not is_session_terminated(exc): + raise + print(f"[MCP] {name}: session terminated by the server; reconnecting") + await self._replace(conn) + except Exception: + # Transport-level failure: the session is suspect, but the call may + # have reached the server, so it is not repeated. + conn.retire() + raise + raise AssertionError("unreachable") + + async def _replace(self, failed: McpConnection) -> None: + async with self._swap_lock: + if self.conn is not failed: + return # a concurrent call already swapped it + if self.send_status: + self.send_status({"type": "status", "message": "Reconnecting to ThoughtSpot..."}) + failed.retire() + self.conn = await mcp_pool.acquire(self.send_status) + self._held.append(self.conn) + +# ThoughtSpot MCP tools, as exposed by the Spotter 3 toolset: +# +# check_connectivity - test connectivity + auth. No inputs. +# search_objects - find existing Liveboards / Answers / Worksheets by name. +# Metadata only; never returns data. +# create_analysis_session - start a session. Optional: data_source_id. +# Returns analytical_session_id. +# send_session_message - ask the Analytics Agent a question. +# Inputs: analytical_session_id, message, additional_context. +# get_session_updates - poll for updates. Input: analytical_session_id. +# Returns session_updates[] + is_done. See AUTOPOLL below. +# create_dashboard - build a dashboard from answer_ids. +# Inputs: title, answers[], note_tile. Returns link. +# list_orgs / switch_org - OAuth-only. The server hides them on `/token/*`, so they +# never appear in the tool list for this static-token setup. +# +# Set ALLOWED_TOOLS to a list of names to restrict the agent. None allows everything +# the server exposes, which is the right default: the tool list is version-negotiated, +# so a hardcoded list silently drops tools added in later API versions. +ALLOWED_TOOLS: list[str] | None = None + +# ── Server-side polling of get_session_updates ────────────────────────────────── +POLL_TOOL = "get_session_updates" +POLL_INITIAL_DELAY = 0.75 # seconds before the first re-poll +POLL_MAX_DELAY = 4.0 # cap on the backoff +POLL_TIMEOUT = 300.0 # give up after this long without is_done + +SYSTEM_PROMPT = """You are a data analyst assistant powered by ThoughtSpot's Analytics Agent. + +Workflow: +- Create one analysis session per conversation with `create_analysis_session`, then ask + questions with `send_session_message`, then call `get_session_updates` once. +- `get_session_updates` is polled to completion for you: a single call returns the Agent's + full response, so never call it twice for the same question. +- Use `search_objects` to find existing Liveboards, Answers or Worksheets by name. It + returns metadata only, never data - to answer a data question, ask the Agent. +- Use `create_dashboard` when the user wants to save or share results, passing the + `answer_id` values from the answers you want on it. + +Presenting answers: +- Every `answer` update is ALREADY rendered in the UI as an interactive ThoughtSpot chart, + in the order it was returned. Do not emit