diff --git a/.claude/launch.json b/.claude/launch.json new file mode 100644 index 0000000..e17e864 --- /dev/null +++ b/.claude/launch.json @@ -0,0 +1,32 @@ +{ + "version": "0.0.1", + "configurations": [ + { + "name": "spotter-backend", + "runtimeExecutable": "bash", + "runtimeArgs": [ + "-c", + "cd mcp/python-react-agent-simple-ui/server && exec ../.venv/bin/python -m uvicorn claude_agent_with_spotter3_mcp_server:app --reload --port 8001" + ], + "port": 8001 + }, + { + "name": "spotter-history-backend", + "runtimeExecutable": "bash", + "runtimeArgs": [ + "-c", + "cd mcp/python-react-agent-simple-ui/server && exec ../.venv/bin/python -m uvicorn claude_agent_with_spotter3_mcp_server_and_chat_history:app --reload --port 8001" + ], + "port": 8001 + }, + { + "name": "spotter-frontend", + "runtimeExecutable": "bash", + "runtimeArgs": [ + "-c", + "cd mcp/python-react-agent-simple-ui/client && exec npm run dev" + ], + "port": 8000 + } + ] +} 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..ea09ad6 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. Both build on the shared agent in `server/spotter3_core.py` (tokens, MCP session, Claude tool loop, SSE), so copy that file along with whichever backend you take: -| 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. +`spotter3_core.py` (shared by both backends) runs 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 (`run_agent()`), 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
; +} + +/** The icon for the mode a click switches to: a sun in dark mode, a moon in light. */ +function ThemeIcon({ theme }: { theme: Theme }) { + return ( + + ); +} + +function App() { + const theme = useTheme(); + 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, theme); + 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..fc3c495 --- /dev/null +++ b/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server.py @@ -0,0 +1,33 @@ +""" +Python agent: Anthropic Claude + the ThoughtSpot MCP server (Spotter 3 toolset). + +FastAPI streams chat responses to the React frontend over Server-Sent Events. The agent +itself - token handling, the shared MCP session, the Claude tool loop, server-side polling +of `get_session_updates` - lives in `spotter3_core.py`. This file is the chat endpoint, +with conversation history held in memory for the life of the process. + +To customise the agent, edit `SYSTEM_PROMPT`, `ALLOWED_TOOLS` and `MODEL` in `spotter3_core.py`. + +MCP server: https://github.com/thoughtspot/mcp-server +""" + +import uuid + +from pydantic import BaseModel + +from spotter3_core import conversations, create_app, run_turn, sse_response + +app = create_app() + + +class ChatRequest(BaseModel): + message: str + response_id: str | None = None + + +@app.post("/api/chat") +async def chat(request: ChatRequest): + conv_id = request.response_id or str(uuid.uuid4()) + print(f"[Chat] conv {conv_id}: {request.message}") + messages = conversations.get(conv_id, []) + [{"role": "user", "content": request.message}] + return sse_response(lambda queue: run_turn(messages, conv_id, queue.put_nowait)) diff --git a/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server_and_chat_history.py b/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server_and_chat_history.py new file mode 100644 index 0000000..ba8987e --- /dev/null +++ b/mcp/python-react-agent-simple-ui/server/claude_agent_with_spotter3_mcp_server_and_chat_history.py @@ -0,0 +1,576 @@ +""" +Python agent: Anthropic Claude + the ThoughtSpot MCP server (Spotter 3 toolset), +with persistent chat history. + +The same agent as `claude_agent_with_spotter3_mcp_server.py` (the agent lives in +`spotter3_core.py`), plus a SQLite chat history, so past conversations survive a restart +and the UI can list, reopen and delete them. It also records the Analytics Agent's work +(`work` events) for the UI's "Show work" section. + +MCP server: https://github.com/thoughtspot/mcp-server +""" + +import asyncio +import json +import os +import sqlite3 +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from urllib.parse import quote + +import httpx2 +from fastapi import HTTPException +from pydantic import BaseModel + +from spotter3_core import ( + TS_HOST, + analytical_sessions, + conversations, + create_app, + run_turn, + server_token, + sse_response, +) + +app = create_app() + +# ════════════════════════════════════════════════════════════════════════════════ +# CHAT HISTORY +# +# Stored in two layers, because the browser and the model need different things: +# +# turns - what the UI renders: user text, assistant text, the answers +# (titles only, see `storable_answers`) and the Agent's work steps. +# claude_messages - the raw Claude message list, tool_use and tool_result blocks +# included. Replayed into the next request so follow-ups keep full +# context after a restart. +# +# The ThoughtSpot MCP server's own conversation storage is internal plumbing for +# `get_session_updates` delivery, not a client-facing history API - so history belongs to +# the app. SQLite via the stdlib, no extra dependency. Calls are small and run in a worker +# thread (`asyncio.to_thread`) to keep the event loop free. +# ════════════════════════════════════════════════════════════════════════════════ + +DB_PATH = Path(os.getenv("CHAT_HISTORY_DB", Path(__file__).resolve().parent / "chat_history.db")) + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS conversations ( + id TEXT PRIMARY KEY, + title TEXT NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + analytical_session_id TEXT, + claude_messages TEXT NOT NULL DEFAULT '[]' +); + +CREATE TABLE IF NOT EXISTS turns ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + role TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + answers TEXT NOT NULL DEFAULT '[]', + work TEXT NOT NULL DEFAULT '[]', + created_at TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS turns_by_conversation ON turns (conversation_id, id); +""" + +TITLE_MAX_LEN = 60 + + +def now_iso() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +@contextmanager +def db() -> Iterator[sqlite3.Connection]: + """A connection that commits on success, rolls back on error and always closes.""" + conn = sqlite3.connect(DB_PATH, timeout=10) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") # concurrent reads while a turn is written + conn.execute("PRAGMA foreign_keys=ON") # so deleting a conversation drops its turns + try: + with conn: + yield conn + finally: + conn.close() + + +def db_init() -> None: + with db() as conn: + conn.executescript(SCHEMA) + # A database from before `work` existed is left untouched by CREATE IF NOT EXISTS. + columns = {row["name"] for row in conn.execute("PRAGMA table_info(turns)")} + if "work" not in columns: + conn.execute("ALTER TABLE turns ADD COLUMN work TEXT NOT NULL DEFAULT '[]'") + print(f"[History] SQLite at {DB_PATH}") + + +db_init() + + +def jsonable(value: Any) -> Any: + """Make a Claude message list JSON-serializable. + + Assistant turns hold SDK block objects. `model_dump` keeps every field the API needs + on replay - including a thinking block's `signature`, which must come back unchanged. + """ + dump = getattr(value, "model_dump", None) + if callable(dump): + return dump(mode="json", exclude_none=True) + if isinstance(value, list): + return [jsonable(item) for item in value] + if isinstance(value, dict): + return {key: jsonable(item) for key, item in value.items()} + return value + + +def storable_answers(answers: list[dict]) -> list[dict]: + """Keep the parts of an answer that survive, drop the parts that go stale. + + An answer object lives about 8 hours on the ThoughtSpot side. After that its + `iframe_url` - and the `answer_id`, really a `{session_id, gen_no}` pair - point at + nothing, and a reopened chat would render a row of error tiles. So only the durable + parts are stored; on replay the client asks the Visual Embed SDK for a live URL from + the conversation id plus the answer's position. `query` is dropped too: nothing reads + it on replay. + + Applied on the way out as well, so rows stored before this rule resolve on replay too. + """ + return [ + {k: v for k, v in answer.items() if k not in ("iframe_url", "answer_id", "query")} + for answer in answers + ] + + +def db_start_conversation(conv_id: str, first_message: str) -> None: + """Create the conversation row if this is its first turn.""" + title = first_message.strip().splitlines()[0][:TITLE_MAX_LEN] or "New chat" + stamp = now_iso() + with db() as conn: + conn.execute( + """ + INSERT INTO conversations (id, title, created_at, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (id) DO UPDATE SET updated_at = excluded.updated_at + """, + (conv_id, title, stamp, stamp), + ) + + +def db_add_turn( + conv_id: str, role: str, content: str, answers: list[dict], work: list[dict] | None = None +) -> None: + stamp = now_iso() + with db() as conn: + conn.execute( + "INSERT INTO turns (conversation_id, role, content, answers, work, created_at)" + " VALUES (?, ?, ?, ?, ?, ?)", + (conv_id, role, content, json.dumps(answers), json.dumps(work or []), stamp), + ) + conn.execute("UPDATE conversations SET updated_at = ? WHERE id = ?", (stamp, conv_id)) + + +def db_save_state( + conv_id: str, claude_messages: list, analytical_session_id: str | None +) -> None: + """Store the raw Claude history and ThoughtSpot session for the next turn.""" + with db() as conn: + conn.execute( + "UPDATE conversations SET claude_messages = ?, analytical_session_id = ?," + " updated_at = ? WHERE id = ?", + (json.dumps(jsonable(claude_messages)), analytical_session_id, now_iso(), conv_id), + ) + + +def db_load_state(conv_id: str) -> tuple[list, str | None]: + with db() as conn: + row = conn.execute( + "SELECT claude_messages, analytical_session_id FROM conversations WHERE id = ?", + (conv_id,), + ).fetchone() + if not row: + return [], None + try: + messages = json.loads(row["claude_messages"]) + except (ValueError, TypeError): + messages = [] + return messages, row["analytical_session_id"] + + +def db_list_conversations(limit: int = 100) -> list[dict]: + with db() as conn: + rows = conn.execute( + """ + SELECT c.id, c.title, c.created_at, c.updated_at, + (SELECT COUNT(*) FROM turns t WHERE t.conversation_id = c.id) AS turn_count + FROM conversations c + WHERE EXISTS (SELECT 1 FROM turns t WHERE t.conversation_id = c.id) + ORDER BY c.updated_at DESC + LIMIT ? + """, + (limit,), + ).fetchall() + return [dict(row) for row in rows] + + +def db_get_conversation(conv_id: str) -> dict | None: + with db() as conn: + row = conn.execute( + "SELECT id, title, created_at, updated_at, analytical_session_id" + " FROM conversations WHERE id = ?", + (conv_id,), + ).fetchone() + if not row: + return None + turns = conn.execute( + "SELECT role, content, answers, work, created_at FROM turns" + " WHERE conversation_id = ? ORDER BY id", + (conv_id,), + ).fetchall() + + return { + **dict(row), + "turns": [ + { + "role": turn["role"], + "content": turn["content"], + "answers": storable_answers(json.loads(turn["answers"] or "[]")), + "work": json.loads(turn["work"] or "[]"), + "created_at": turn["created_at"], + } + for turn in turns + ], + } + + +def db_get_session_id(conv_id: str) -> str | None: + with db() as conn: + row = conn.execute( + "SELECT analytical_session_id FROM conversations WHERE id = ?", (conv_id,) + ).fetchone() + return row["analytical_session_id"] if row else None + + +def db_delete_conversation(conv_id: str) -> bool: + with db() as conn: + return conn.execute("DELETE FROM conversations WHERE id = ?", (conv_id,)).rowcount > 0 + + +def db_rename_conversation(conv_id: str, title: str) -> bool: + with db() as conn: + updated = conn.execute( + "UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?", + (title.strip()[:TITLE_MAX_LEN] or "New chat", now_iso(), conv_id), + ).rowcount + return updated > 0 + + +async def load_conversation_state(conv_id: str) -> list: + """Claude message history for a conversation: from the in-memory cache, else SQLite + (e.g. after a restart).""" + if conv_id not in conversations: + messages, session_id = await asyncio.to_thread(db_load_state, conv_id) + conversations[conv_id] = messages + if session_id: + analytical_sessions[conv_id] = session_id + return conversations[conv_id] + + +class StreamRecorder: + """Fans every SSE event out to the browser and to the transcript we persist. + + The UI renders a turn from exactly the events it received, so recording them here is + what makes a reopened conversation look like the live one. + """ + + def __init__(self, queue: asyncio.Queue) -> None: + self.queue = queue + self.text_parts: list[str] = [] + self.answers: list[dict] = [] + self.work: list[dict] = [] # the Analytics Agent's steps, for "Show work" + + def send(self, event: dict) -> None: + kind = event.get("type") + if kind == "delta": + self.text_parts.append(event.get("text") or "") + elif kind == "answer": + self.answers.append({k: event.get(k) for k in ("answer_id", "title", "iframe_url")}) + elif kind == "work": + append_work(self.work, event) + self.queue.put_nowait(event) + + @property + def text(self) -> str: + return "".join(self.text_parts) + + +def append_work(work: list[dict], event: dict) -> None: + """Add one `work` event to a turn's work list, the way the client does. + + Reasoning prose streams a few words per event, so a `thought` continues the previous + one instead of starting a new step. The client applies the same rule, which keeps the + live and replayed views identical. + """ + kind = event.get("kind") + text = event.get("text") or "" + if kind == "thought" and work and work[-1].get("kind") == "thought": + work[-1]["text"] += text + return + step = {"kind": kind, "text": text} + if event.get("query"): + step["query"] = event["query"] + work.append(step) + + +async def persist_turn(conv_id: str, recorder: StreamRecorder, history: list | None) -> None: + """Store the assistant turn, and the state the next turn needs. + + Also runs when a turn fails or is cancelled: the user turn is already stored, so + skipping this would leave a question with no visible answer on reopen. + """ + if recorder.text or recorder.answers or recorder.work: + await asyncio.to_thread( + db_add_turn, + conv_id, + "assistant", + recorder.text, + storable_answers(recorder.answers), + recorder.work, + ) + if history is not None: + await asyncio.to_thread( + db_save_state, conv_id, history, analytical_sessions.get(conv_id) + ) + + +class ChatRequest(BaseModel): + message: str + response_id: str | None = None + + +class RenameRequest(BaseModel): + title: str + + +@app.post("/api/chat") +async def chat(request: ChatRequest): + conv_id = request.response_id or str(uuid.uuid4()) + print(f"[Chat] conv {conv_id}: {request.message}") + + history = await load_conversation_state(conv_id) + messages = history + [{"role": "user", "content": request.message}] + + await asyncio.to_thread(db_start_conversation, conv_id, request.message) + await asyncio.to_thread(db_add_turn, conv_id, "user", request.message, []) + + async def turn(queue: asyncio.Queue) -> None: + recorder = StreamRecorder(queue) + await run_turn( + messages, + conv_id, + recorder.send, + persist=lambda history: persist_turn(conv_id, recorder, history), + show_work=True, + ) + + return sse_response(turn) + + +# ── ThoughtSpot conversation lookups ──────────────────────────────────────────── +# A stored conversation is checked against what ThoughtSpot holds, so replayed charts +# line up with the answers the SDK resolves. + + +async def ts_request( + method: str, + path: str, + label: str, + read_timeout: float, + headers: dict | None = None, + json_body: dict | None = None, +) -> Any: + """Call the ThoughtSpot REST API as the server user. Returns the JSON body, or None + (logged) when the call fails.""" + try: + async with httpx2.AsyncClient(timeout=httpx2.Timeout(10.0, read=read_timeout)) as client: + response = await client.request( + method, + f"{TS_HOST.rstrip('/')}{path}", + headers={ + "Authorization": f"Bearer {await server_token()}", + "Accept": "application/json", + **(headers or {}), + }, + json=json_body, + ) + if response.status_code != 200: + print(f"[History] {label} -> {response.status_code}") + return None + return response.json() + except (httpx2.HTTPError, ValueError, TypeError, AttributeError) as exc: + print(f"[History] {label} failed: {type(exc).__name__}: {exc}") + return None + + +async def ts_conversation_messages(session_id: str) -> list[dict] | None: + """A ThoughtSpot conversation's messages, oldest first, or None if unreadable.""" + body = await ts_request( + "GET", + f"/api/rest/2.0/ai/agent/conversations/{session_id}/messages", + label=f"getConversation {session_id}", + read_timeout=45.0, + ) + if not isinstance(body, dict): + return None + messages = body.get("messages") + return messages if isinstance(messages, list) else [] + + +def answer_items(message: dict, thinking: bool) -> list[dict]: + """A message's answer items that are (or are not) the Agent's thinking answers. + + Compared with `is` on purpose: an item missing `is_thinking` counts as neither. Items + without an `answer_id` are skipped too. Both are the rules the Visual Embed SDK applies + when it replays stored answers, so our counts match its indexes. + """ + items = message.get("response_items") if isinstance(message, dict) else None + return [ + item + for item in (items or []) + if isinstance(item, dict) + and item.get("type") == "answer" + and item.get("is_thinking") is thinking + and item.get("answer_id") + ] + + +async def ts_load_answer(session_id: str, answer_id: str) -> dict | None: + """Live embed ids for one answer of a conversation, or None if it cannot load. + + The answer behind a stored URL expires after about 8 hours; loading it again through + the conversation service gives fresh ids for the same answer. This is the call the + Visual Embed SDK makes when it replays a stored answer. + """ + body = await ts_request( + "POST", + f"/conversation/v2/{session_id}/message/{quote(answer_id, safe='')}/load/public", + label=f"load answer {answer_id}", + read_timeout=120.0, # loading re-runs the query; 45s+ is common on a loaded cluster + headers={"Content-Type": "application/json", "x-requested-by": "ThoughtSpot"}, + json_body={"type": "TS_ANSWER"}, + ) + if not isinstance(body, dict): + return None + + answer = body.get("answer") or {} + ac_state = answer.get("ac_state") or {} + params = { + "session_id": answer.get("session_identifier"), + "gen_no": answer.get("generation_number"), + "ac_session_id": ac_state.get("transaction_identifier"), + "ac_gen_no": ac_state.get("generation_number"), + } + # All four are needed to build the embed route; a partial set renders an error. + return params if all(params.values()) else None + + +def reconcile_answers(conversation: dict, ts_messages: list[dict] | None) -> dict: + """Align a stored conversation with the answers ThoughtSpot actually holds. + + The stored copy is not a reliable source of truth. A stream can be cut off (the + browser navigates away, the process restarts) after the Agent was already asked; it + finishes anyway, so the answer exists on ThoughtSpot's side while our turn recorded + none of it. And `answer_index` has to count the way ThoughtSpot counts, or a replayed + answer resolves to the wrong chart. + + So ThoughtSpot decides how many answers each turn has and the stored rows only supply + titles. A turn missing answers gets untitled placeholders; the client renders them and + the SDK resolves each by its index. + """ + turns = conversation.get("turns") or [] + if ts_messages is None: + # ThoughtSpot unreachable: fall back to the stored shape rather than drop charts. + index = 0 + for turn in turns: + for answer in turn.get("answers") or []: + answer["answer_index"] = index + index += 1 + return conversation + + # One ThoughtSpot message per user prompt: pair them with the assistant turns in order. + assistant_turns = [turn for turn in turns if turn.get("role") == "assistant"] + index = 0 + for position, turn in enumerate(assistant_turns): + message = ts_messages[position] if position < len(ts_messages) else {} + expected = len(answer_items(message, thinking=False)) + # The thinking answers behind this turn's Show work queries, in step order. The + # client loads one by id when its row is expanded. + thinking_answers = answer_items(message, thinking=True) + turn["work_answer_ids"] = [item["answer_id"] for item in thinking_answers] + answers = turn.get("answers") or [] + # Titles we recorded, padded out to the count ThoughtSpot reports. + merged = answers[:expected] + [{} for _ in range(max(0, expected - len(answers)))] + for answer in merged: + answer["answer_index"] = index + index += 1 + turn["answers"] = merged + return conversation + + +# ── Chat history endpoints ────────────────────────────────────────────────────── + + +@app.get("/api/conversations") +async def list_conversations(): + return {"conversations": await asyncio.to_thread(db_list_conversations)} + + +@app.get("/api/conversations/{conv_id}") +async def get_conversation(conv_id: str): + conversation = await asyncio.to_thread(db_get_conversation, conv_id) + if not conversation: + raise HTTPException(status_code=404, detail="Conversation not found") + + session_id = conversation.get("analytical_session_id") + ts_messages = await ts_conversation_messages(session_id) if session_id else None + return reconcile_answers(conversation, ts_messages) + + +@app.get("/api/conversations/{conv_id}/work-answers/{answer_id}") +async def get_work_answer(conv_id: str, answer_id: str): + """Live embed ids for a thinking answer shown in a stored turn's Show work. + + `answer_id` comes from the turn's `work_answer_ids`. The client asks only when a step + is expanded, so a reopened chat does not load every intermediate chart up front. + ThoughtSpot checks the answer belongs to the conversation. + """ + session_id = await asyncio.to_thread(db_get_session_id, conv_id) + if not session_id: + raise HTTPException(status_code=404, detail="Conversation has no ThoughtSpot session") + + frame_params = await ts_load_answer(session_id, answer_id) + if not frame_params: + raise HTTPException(status_code=502, detail="Could not load that answer") + return {"frame_params": frame_params} + + +@app.patch("/api/conversations/{conv_id}") +async def rename_conversation(conv_id: str, request: RenameRequest): + if not await asyncio.to_thread(db_rename_conversation, conv_id, request.title): + raise HTTPException(status_code=404, detail="Conversation not found") + return {"status": "ok"} + + +@app.delete("/api/conversations/{conv_id}") +async def delete_conversation(conv_id: str): + if not await asyncio.to_thread(db_delete_conversation, conv_id): + raise HTTPException(status_code=404, detail="Conversation not found") + conversations.pop(conv_id, None) + analytical_sessions.pop(conv_id, None) + return {"status": "deleted"} diff --git a/mcp/python-react-agent-simple-ui/server/requirements.txt b/mcp/python-react-agent-simple-ui/server/requirements.txt index 4da0eef..789e0be 100644 --- a/mcp/python-react-agent-simple-ui/server/requirements.txt +++ b/mcp/python-react-agent-simple-ui/server/requirements.txt @@ -1,7 +1,7 @@ openai>=2.4.0 -anthropic>=0.52.0 -fastapi>=0.115.0 -uvicorn>=0.34.0 +anthropic>=1.2.0,<2 +fastapi>=0.141.0 +uvicorn>=0.52.0 python-dotenv>=1.1.1 -mcp>=1.9.0 -httpx>=0.27.0 +mcp>=2.1.1,<3 +httpx2>=2.0.0 diff --git a/mcp/python-react-agent-simple-ui/server/spotter3_core.py b/mcp/python-react-agent-simple-ui/server/spotter3_core.py new file mode 100644 index 0000000..040733a --- /dev/null +++ b/mcp/python-react-agent-simple-ui/server/spotter3_core.py @@ -0,0 +1,1061 @@ +""" +Shared core of the Claude + ThoughtSpot MCP server (Spotter 3 toolset) examples. + +Both `claude_agent_with_spotter3_mcp_server.py` and +`claude_agent_with_spotter3_mcp_server_and_chat_history.py` build on this module: token +handling, the shared MCP session, the Claude tool loop and the SSE plumbing. + +Why client-side MCP: the ThoughtSpot MCP server's static-token endpoint needs custom +HTTP headers (Authorization + x-ts-host), which Anthropic's server-side MCP connector +cannot send. So this process connects to the MCP server itself and runs the tool calls. + +Two things it does that a plain pass-through loop does not: + +1. It polls `get_session_updates` itself. The Analytics Agent answers asynchronously, so + one call usually returns `is_done: false`. Letting the model poll costs a model + round-trip per poll; here we poll until done and hand the model one consolidated + result, streaming the Agent's progress to the UI as it arrives. + +2. The client renders answers, not the model. Each `answer` update becomes an `answer` + event; the React client mounts an iframe for it, which the Visual Embed SDK's + `startAutoMCPFrameRenderer` upgrades into a real ThoughtSpot embed. The URL stays out + of what the model sees. + +MCP server: https://github.com/thoughtspot/mcp-server +""" + +import asyncio +import json +import os +import time +import traceback +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, 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 APIRouter, 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 + +load_dotenv(dotenv_path=Path(__file__).resolve().parent.parent / ".env") + +Send = Callable[[dict], None] # pushes one SSE event to the browser + +# ── Claude ────────────────────────────────────────────────────────────────────── +claude_client = anthropic.AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + +# Haiku 4.5: the agent only picks tools, passes the question on and summarises - the +# Analytics Agent does the analysis. Measured on one question: ~4s of model time on +# Haiku, ~17s on Sonnet 5, ~8s on Opus 5. Set ANTHROPIC_MODEL=claude-sonnet-5 for +# stronger reasoning on multi-step follow-ups. +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 it to another model, but only models with published +# fallback targets accept it. Sonnet 5 and Haiku 4.5 have none, so a refusal is reported. +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, and picking tools does not need it. + 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") + +# `/token/*` is the static-bearer-token transport (`/bearer/*` is legacy and 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. Append `&enable-raw-session-updates=true` to TS_MCP_URL to +# stream the Agent's own updates; both shapes are read (see "Session updates"). +MCP_API_VERSION = os.getenv("TS_MCP_API_VERSION", "latest") +MCP_URL = os.getenv( + "TS_MCP_URL", + f"https://agent.thoughtspot.app/token/mcp?api-version={MCP_API_VERSION}", +) + +# ── ThoughtSpot tokens ────────────────────────────────────────────────────────── +# Two consumers, two policies: +# +# The browser. The Visual Embed SDK needs a FRESH token on every `getAuthToken` call, +# 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 then every tool call fails ("Failed to validate connection") while the +# tool *list* still works. So we mint a token and cache it until shortly before it +# expires. TS_AUTH_TOKEN is only the fallback when no minting credentials are set. +# +# Minting needs TS_EMBED_USERNAME plus the cluster secret key (Develop > Customizations > +# Security Settings > Trusted authentication) or that user's password. The secret key wins +# 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")) +# The server's own token lives longer: it is never handed out, and fewer mints are +# easier on a slow cluster. +SERVER_TOKEN_VALIDITY_SEC = int(os.getenv("TS_SERVER_TOKEN_VALIDITY_SEC", "3600")) + +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." + ) + +# A cached token is treated as expired this long before it really is, so a call already +# in flight does not race the expiry. +TOKEN_REFRESH_MARGIN_SEC = 120 +# The background task renews this long before expiry (capped at half the token's life), so +# no request waits on a mint - which can take 40s on a loaded cluster. +TOKEN_BACKGROUND_LEAD_SEC = 720 +TOKEN_BACKGROUND_RETRY_SEC = 30.0 # after a failed background mint; the old token keeps serving +# The loop re-reads the wall clock this often: asyncio.sleep pauses with a sleeping laptop. +TOKEN_BACKGROUND_TICK_SEC = 60.0 + +_server_token: str | None = None +# Wall-clock (time.time()) instants - not time.monotonic(), which stops while macOS sleeps +# and would keep serving a token the cluster has already expired. +_server_token_expiry = 0.0 # stop serving the token +_server_token_refresh_at = 0.0 # background task renews +# One mint at a time: an agent turn fires several tool calls at once, and on a cold cache +# they would each mint a token. +_server_token_lock = asyncio.Lock() + + +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} + if TS_SECRET_KEY: + payload["secret_key"] = TS_SECRET_KEY + else: + payload["password"] = TS_EMBED_PASSWORD + + # The endpoint is slow at times (15s is normal on a loaded cluster); catch the + # timeout so the browser gets an explanation instead of a bare 500. + 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(type(exc).__name__, 504) from exc + + if response.status_code != 200: + # The cluster's own reason (bad secret key, trusted auth disabled, ...). + 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. + 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: + # An expiring token 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.""" + 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 + + +# ── MCP connection ────────────────────────────────────────────────────────────── +# 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 ("MCPError: +# Server returned an error response") that a second attempt gets past. So a failed +# handshake is retried once, with a fresh token to rule 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 status(message: str) -> dict: + return {"type": "status", "message": message} + + +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: Send | None = None) -> AsyncIterator[tuple[Client, Any, str]]: + """An initialised MCP session, its tool listing and the bearer token it runs on.""" + for attempt in range(1, MCP_CONNECT_ATTEMPTS + 1): + stack = AsyncExitStack() + try: + token = await server_token() + headers = {"Authorization": f"Bearer {token}", "x-ts-host": TS_HOST} + 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: + send(status("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 a session costs ~10s on a loaded cluster (discover probe, initialize, +# tools/list), so one session is shared by every turn and reopened only when: +# +# - 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). +# +# A background task opens and closes the session, never a request: anyio requires the +# task that opened a session to close it, and closing it in a request task got cancelled +# half-way whenever the browser stream ended. + +# The error the MCP client raises once the server forgets the session. Only this pairing +# means the call never ran and is safe to repeat: -32600 alone 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() + + # 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 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. + self._lock = asyncio.Lock() + self._conn: McpConnection | None = None + + async def acquire(self, send: Send | 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) + conn.leases += 1 + self._conn = conn + return conn + + async def _open(self, send: Send | None) -> McpConnection: + ready: asyncio.Future = asyncio.get_running_loop().create_future() + stop = asyncio.Event() + + async def own() -> None: + try: + async with connect_mcp(send) 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: + stop.set() # close the session once it opens rather than leave it ownerless + if not asyncio.current_task().cancelling(): + # Not this turn that was cancelled but the handshake (e.g. at shutdown): + # report a failure instead of 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: Send | None = None) -> None: + self.send = send + 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) + 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: + self.send(status("Reconnecting to ThoughtSpot...")) + failed.retire() + self.conn = await mcp_pool.acquire(self.send) + self._held.append(self.conn) + + +# ── Tools and prompt ──────────────────────────────────────────────────────────── +# The Spotter 3 toolset: check_connectivity, search_objects (metadata only, never data), +# create_analysis_session, send_session_message, get_session_updates, create_dashboard. +# `list_orgs` / `switch_org` are OAuth-only and never appear on `/token/*`. +# +# 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 + +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 + +TOOL_STATUS = { + "check_connectivity": "Checking the ThoughtSpot connection...", + "search_objects": "Searching ThoughtSpot...", + "create_analysis_session": "Starting an analysis session...", + "send_session_message": "Asking the Analytics Agent...", + POLL_TOOL: "Waiting for the Analytics Agent...", + "create_dashboard": "Building the dashboard...", +} + +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