diff --git a/AGENTS.md b/AGENTS.md index 0596a8f5b..4ffaf3509 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -30,6 +30,14 @@ For detailed instructions on writing tests (mocking patterns, test structure, MC - **Main skill**: `.github/skills/test-writing/SKILL.md` - **MCP-specific patterns**: `.github/skills/test-writing/mcp-tools-testing.md` +## Linting + +### Running ESLint + +- Use `npm run test:lint` to check all TypeScript files for linting errors +- **IMPORTANT**: The `get_errors` tool may not catch all linting issues, especially in files that aren't currently open in the editor +- Always run `npm run test:lint` before committing to ensure all files pass linting + ## MCP Tools ### Writing MCP Tools diff --git a/docs/mcp-session-architecture.md b/docs/mcp-session-architecture.md new file mode 100644 index 000000000..2cd3ce339 --- /dev/null +++ b/docs/mcp-session-architecture.md @@ -0,0 +1,337 @@ +# MCP Session Architecture + +## Understanding the Layers + +There are distinct layers in the MCP server architecture: + +### Layer 1: HTTP Server (Network Layer) + +- **ONE** `http.Server` instance per VS Code extension +- Listens on a **single port** (e.g., `http://localhost:45678/mcp`) +- Receives all incoming HTTP requests +- Routes requests based on session ID +- Lives for the entire lifetime of the extension + +### Layer 2: MCP SDK Server Instances (Protocol Layer) + +- Handles MCP protocol logic (tool registration, request/response) +- **Current (Stateless)**: ONE shared instance +- **New (Stateful)**: MULTIPLE instances (one per session) + +### Layer 3: Transport Layer + +- Manages request/response streaming +- **Current (Stateless)**: New transport per HTTP request +- **New (Stateful)**: One transport per session, reused across requests + +## Current Architecture (Stateless) + +``` +┌─────────────────────────────────────────────────────┐ +│ VS Code Extension Process │ +│ │ +│ ┌────────────────────────────────────────────┐ │ +│ │ McpServer Class Instance │ │ +│ │ │ │ +│ │ http.Server (Port 45678) ◄───────────────┼─────┼─── Client Request 1 +│ │ │ │ │ +│ │ │ │ │ +│ │ ├─► Create Transport 1 │ │ +│ │ │ Connect SdkMcpServer ──────┐ │ │ +│ │ │ Handle Request │ │ │ +│ │ │ Close Transport │ │ │ +│ │ │ │ │ │ +│ │ ├─► Create Transport 2 ◄──────┼───┼─────┼─── Client Request 2 +│ │ │ Connect SdkMcpServer ──┐ │ │ │ (parallel) +│ │ │ Handle Request │ │ │ │ +│ │ │ Close Transport │ │ │ │ +│ │ │ │ │ │ │ +│ │ SHARED SdkMcpServer Instance ◄───┴───┴───┼─────┼─── ⚠️ RACE CONDITION! +│ │ (ONE instance, multiple connections) │ │ +│ │ │ │ +│ └────────────────────────────────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────┘ +``` + +**Problem**: Multiple transports trying to connect to the **same** SDK server instance simultaneously. + +## New Architecture (Stateful with Sessions) + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ VS Code Extension Process │ +│ │ +│ ┌──────────────────────────────────────────────────────────┐ │ +│ │ McpServer Class Instance │ │ +│ │ │ │ +│ │ http.Server (Port 45678) ◄─────────────────────────────┼──┼─── Initialize Request +│ │ │ │ │ (no session ID) +│ │ │ │ │ +│ │ ├─► Detect Initialize Request │ │ +│ │ │ Create NEW SdkMcpServer for Session A │ │ +│ │ │ Create Transport A │ │ +│ │ │ Connect Server A to Transport A (ONCE) │ │ +│ │ │ Store in Maps │ │ +│ │ │ Return session ID: "session-abc" │ │ +│ │ │ │ │ +│ │ ├─► Request with session-abc ◄──────────────────┼──┼─── Follow-up Request 1 +│ │ │ Lookup Transport A │ │ +│ │ │ Reuse (no new connection) │ │ +│ │ │ │ │ +│ │ ├─► Request with session-abc ◄──────────────────┼──┼─── Follow-up Request 2 +│ │ │ Lookup Transport A │ │ (parallel) +│ │ │ Reuse (no new connection) │ │ +│ │ │ │ │ +│ │ ├─► Initialize Request (different client) ◄──────┼──┼─── New Session +│ │ │ Create NEW SdkMcpServer for Session B │ │ +│ │ │ Create Transport B │ │ +│ │ │ Connect Server B to Transport B (ONCE) │ │ +│ │ │ Store in Maps │ │ +│ │ │ Return session ID: "session-xyz" │ │ +│ │ │ │ │ +│ │ Session Storage: │ │ +│ │ ┌─────────────────────────────────────────────┐ │ │ +│ │ │ transports Map │ │ │ +│ │ │ "session-abc" → Transport A ──┐ │ │ │ +│ │ │ "session-xyz" → Transport B ──┼─┐ │ │ │ +│ │ └──────────────────────────────────┼─┼─────────┘ │ │ +│ │ │ │ │ │ +│ │ ┌─────────────────────────────────┼─┼─────────┐ │ │ +│ │ │ servers Map │ │ │ │ │ +│ │ │ "session-abc" → SdkMcpServer A│ │ │ │ │ +│ │ │ "session-xyz" → SdkMcpServer B │ │ │ │ +│ │ └────────────────────────────────────┼─────────┘ │ │ +│ │ │ │ │ +│ │ SdkMcpServer Instance A ◄───────────┘ │ │ +│ │ (tools registered, connected to Transport A) │ │ +│ │ │ │ +│ │ SdkMcpServer Instance B ◄──────────────────────────┐ │ │ +│ │ (tools registered, connected to Transport B) │ │ │ +│ │ │ │ │ +│ └───────────────────────────────────────────────────────┘ │ +│ │ +└──────────────────────────────────────────────────────────────┘ +``` + +**Solution**: Each session has its own isolated SDK server + transport pair. + +## Key Points + +### 1. HTTP Server (Always ONE) + +```typescript +class McpServer { + private httpServer: http.Server | null = null; // ← ONE instance + + async start(port: number) { + // Create ONE HTTP server that listens on ONE port + this.httpServer = http.createServer(async (req, res) => { + // Route to appropriate session based on session ID + }); + + this.httpServer.listen(port); + } +} +``` + +- **Never changes**: Always one HTTP server per extension instance +- **Port**: Single port shared by all sessions +- **Routing**: Uses `mcp-session-id` header to route to correct session + +### 2. SDK Server Instances (ONE → MANY) + +**Current (Stateless)**: + +```typescript +class McpServer { + private server: SdkMcpServer; // ← Shared by all requests + + constructor() { + this.server = new SdkMcpServer({...}); + // Register tools once + } +} +``` + +**New (Stateful)**: + +```typescript +class McpServer { + private servers: Map; // ← One per session + + private createServer(): SdkMcpServer { + const server = new SdkMcpServer({...}); + // Register tools on this instance + return server; + } + + async handleInitialize() { + const server = this.createServer(); // New instance! + const sessionId = randomUUID(); + this.servers.set(sessionId, server); + } +} +``` + +### 3. Request Flow Examples + +#### Example 1: Single Client, Multiple Requests + +``` +Client 1 (VS Code Copilot) + │ + ├─► POST /mcp (initialize) ──┐ + │ No session ID │ + │ ← Response: session-abc │ + │ │ Same Session + ├─► POST /mcp (list tools) ──┤ Same Server Instance + │ Session: session-abc │ Same Transport + │ │ + ├─► POST /mcp (call tool) ──┤ + │ Session: session-abc │ + │ │ + └─► POST /mcp (call tool) ──┘ + Session: session-abc +``` + +**Result**: + +- 4 HTTP requests → ONE HTTP server +- 1 session → ONE SDK server instance +- 1 session → ONE transport (reused 4 times) + +#### Example 2: Multiple Clients (Parallel Sessions) + +``` +Client 1 (VS Code) Client 2 (Windsurf) + │ │ + ├─► POST /mcp (init) ├─► POST /mcp (init) + │ ← session-abc │ ← session-xyz + │ │ + │ Different Sessions │ + │ Different SDK Servers │ + │ Isolated from each other │ + │ │ + ├─► POST /mcp (tool) ├─► POST /mcp (tool) + │ session-abc │ session-xyz + │ │ + └─► POST /mcp (tool) └─► POST /mcp (tool) + session-abc session-xyz +``` + +**Result**: + +- 6 HTTP requests → ONE HTTP server (handles all) +- 2 sessions → TWO SDK server instances +- 2 transports (one per session) + +#### Example 3: Parallel Requests in Same Session + +``` +Client (parallel tool calls) + │ + ├──┬─► POST /mcp (tool A) ─┐ + │ │ session-abc │ + │ │ ├─► Same Transport + │ └─► POST /mcp (tool B) ─┘ Queued internally + │ session-abc + │ (parallel) +``` + +**Result**: + +- Both requests arrive at HTTP server simultaneously +- Both lookup same transport from `transports.get("session-abc")` +- Transport handles queueing internally +- No race condition because server is already connected + +## Why This Architecture? + +### Network Constraints + +- **ONE port per service**: Can't have multiple HTTP servers on same port +- **Solution**: One HTTP server routes to multiple sessions + +### Protocol Isolation + +- **Sessions must be isolated**: Different clients shouldn't interfere +- **Solution**: Separate SDK server instance per session + +### Connection Stability + +- **Avoid reconnection overhead**: `server.connect()` should happen once +- **Solution**: Create connection during initialization, reuse transport + +## Memory Implications + +### Current (Stateless) + +``` +Memory per request: +- Transport: ~1KB +- Connection overhead: ~100ms +- Total: Minimal but inefficient (created/destroyed constantly) +``` + +### New (Stateful) + +``` +Memory per session: +- SDK Server instance: ~50KB +- Transport: ~1KB +- Event store (optional): ~10KB +- Total: ~60KB per active session + +Typical usage: +- 1-2 active sessions (one per IDE) +- ~120KB total +- Sessions cleaned up when idle +``` + +**Trade-off**: Slightly more memory for much better performance and new capabilities. + +## Implementation Details + +### HTTP Request Handler Pseudocode + +```typescript +this.httpServer = http.createServer(async (req, res) => { + // Extract session ID from header + const sessionId = req.headers["mcp-session-id"]; + + if (sessionId && this.transports.has(sessionId)) { + // EXISTING SESSION: Reuse transport + const transport = this.transports.get(sessionId); + await transport.handleRequest(req, res, body); + } else if (!sessionId && isInitializeRequest(body)) { + // NEW SESSION: Create server + transport + const server = this.createServer(); // New SdkMcpServer + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: sessionId => { + this.servers.set(sessionId, server); + this.transports.set(sessionId, transport); + }, + }); + + await server.connect(transport); // Connect ONCE + await transport.handleRequest(req, res, body); + } else { + // ERROR: Invalid request + res.status(400).json({ error: "Invalid session" }); + } +}); +``` + +## Summary + +| Layer | Current (Stateless) | New (Stateful) | +| --------------- | ------------------- | --------------------------- | +| **HTTP Server** | 1 instance | 1 instance (no change) | +| **Port** | 1 port | 1 port (no change) | +| **SDK Servers** | 1 shared instance | N instances (1 per session) | +| **Transports** | 1 per request | 1 per session | +| **Connections** | N per request | 1 per session | + +The HTTP server is just a **router** - it receives requests and dispatches them to the appropriate session's SDK server instance based on the session ID. diff --git a/docs/mcp.md b/docs/mcp.md index 273ca8b9e..8cf2bf342 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -141,7 +141,7 @@ The MCP server provides tools for: - `getLogs` - Retrieve server or debug logs. - `showOutputPanel` - Display output panel in VS Code. -For detailed documentation on each tool including parameters, return types, and examples, see [MCP Tool Reference](mcp-tools.md). +For detailed documentation on each tool including parameters, return types, and examples, see the [MCP Tool Reference](mcp-tools.md). For technical details about the MCP server's session-based architecture, see the [MCP Session Architecture](mcp-session-architecture.md) reference. ## Chat Skills @@ -174,6 +174,46 @@ The documentation server can be independently enabled/disabled via the `deephave When `deephaven.mcp.enabled` is `true`, documentation queries are enabled by default. Set `deephaven.mcp.docsEnabled` to `false` to disable documentation queries while keeping the extension's MCP tools available. +## Troubleshooting + +### Session Not Found (404) + +**Symptom**: The MCP server returns a `404` error with a message about the session not being found. + +**Causes**: + +- The session expired or was cleaned up (e.g., after extension restart or VS Code reload). +- The client is sending a stale session ID from a previous connection. + +**Resolution**: Restart the AI assistant session or MCP client. Most clients will automatically re-initialize and obtain a new session ID. + +### Session Initialization Failures (400) + +**Symptom**: Requests fail with a `400 Bad Request` error. + +**Causes**: + +- A request was sent without a session ID but was not an `initialize` request. +- The client is not following the MCP session protocol. + +**Resolution**: Ensure the client sends an `initialize` request first to establish a session before sending other requests. + +### Stale Sessions After Extension Restart + +**Symptom**: MCP tools stop responding or return errors after the Deephaven extension restarts. + +**Cause**: When the extension restarts, all active sessions are terminated. Clients holding old session IDs will receive errors. + +**Resolution**: Restart the AI assistant session or reload the MCP configuration so the client re-initializes with a fresh session. + +### Port Changes After Workspace Switch + +**Symptom**: MCP connection fails after switching VS Code workspaces. + +**Cause**: Each workspace uses an auto-allocated port. Switching workspaces changes the port. + +**Resolution**: Check the `MCP:` status bar item for the current port and update your MCP configuration accordingly. You may also need to restart the AI assistant session. + ## Tool Response Format All MCP tools follow a consistent response structure: diff --git a/package-lock.json b/package-lock.json index 90a4a854f..5abbce9c0 100644 --- a/package-lock.json +++ b/package-lock.json @@ -14,6 +14,7 @@ "@deephaven-enterprise/auth-nodejs": "^1.20250219.134-beta", "@deephaven-enterprise/query-utils": "^1.20250219.134-beta", "@deephaven/jsapi-nodejs": "^1.15.0", + "@deephaven/jsapi-utils": "^1.16.0", "@modelcontextprotocol/sdk": "^1.27.1", "archiver": "^7.0.1", "chai": "^4.5.0", @@ -24,10 +25,12 @@ "@cfworker/json-schema": "^4.1.1", "@deephaven-enterprise/jsapi-types": "^1.20250219.134-beta", "@deephaven/jsapi-types": "^41.2.0", + "@react-types/shared": "^3.33.1", "@types/archiver": "^6.0.3", "@types/chai": "^4.3.20", "@types/mocha": "^10.0.10", "@types/node": "22.19.1", + "@types/react": "^19.2.14", "@types/vscode": "^1.105.0", "@types/vscode-webview": "^1.57.5", "@types/ws": "^8.5.10", @@ -600,6 +603,15 @@ "integrity": "sha512-NwMxFmNCnRV4/2+MdN/8vUGiEtXFgL1K/+iXTKKvi+Brje5JHOSCn2miCKR9tAn0LNb/UdmJq+DSIZqvz8cU/Q==", "license": "Apache-2.0" }, + "node_modules/@deephaven/filters": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@deephaven/filters/-/filters-1.1.0.tgz", + "integrity": "sha512-SMWKYT8aFtZ/CyVUtUiF1/RPQGl4Y6dvNG43KvmIMKKMXYqhliKD291aynYSl7C8IafkHNzdDtdAZPgPXLoOWA==", + "license": "Apache-2.0", + "engines": { + "node": ">=16" + } + }, "node_modules/@deephaven/jsapi-nodejs": { "version": "1.15.0", "resolved": "https://registry.npmjs.org/@deephaven/jsapi-nodejs/-/jsapi-nodejs-1.15.0.tgz", @@ -655,6 +667,56 @@ "dev": true, "license": "Apache-2.0" }, + "node_modules/@deephaven/jsapi-utils": { + "version": "1.16.0", + "resolved": "https://registry.npmjs.org/@deephaven/jsapi-utils/-/jsapi-utils-1.16.0.tgz", + "integrity": "sha512-lyMDgxmb7v2GDlm1Ksf4S037QuWUhpOMKtnHZoEuKKIthcd6WejsNP+S8U2iRCYSJiJ1jEpTAVMrrS2CM4ZC7w==", + "license": "Apache-2.0", + "dependencies": { + "@deephaven/filters": "^1.1.0", + "@deephaven/jsapi-types": "^1.0.0-dev0.40.4", + "@deephaven/log": "^1.8.0", + "@deephaven/utils": "^1.10.0", + "lodash.clamp": "^4.0.3", + "nanoid": "^5.0.7" + }, + "engines": { + "node": ">=16" + } + }, + "node_modules/@deephaven/jsapi-utils/node_modules/@deephaven/jsapi-types": { + "version": "1.0.0-dev0.40.9", + "resolved": "https://registry.npmjs.org/@deephaven/jsapi-types/-/jsapi-types-1.0.0-dev0.40.9.tgz", + "integrity": "sha512-NwMxFmNCnRV4/2+MdN/8vUGiEtXFgL1K/+iXTKKvi+Brje5JHOSCn2miCKR9tAn0LNb/UdmJq+DSIZqvz8cU/Q==", + "license": "Apache-2.0" + }, + "node_modules/@deephaven/jsapi-utils/node_modules/@deephaven/log": { + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@deephaven/log/-/log-1.8.0.tgz", + "integrity": "sha512-gzp6/7qW4W8Re+DLSaG33KEQJ30OrrNq3cNQA8fUeXQrabSNOIsyeVOaerQ/57c4zWhWVKamplax0LIYRsmDiw==", + "license": "Apache-2.0", + "dependencies": { + "event-target-shim": "^6.0.2", + "jszip": "^3.10.1", + "safe-stable-stringify": "^2.5.0" + }, + "engines": { + "node": ">=16" + } + }, + "node_modules/@deephaven/jsapi-utils/node_modules/@deephaven/utils": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@deephaven/utils/-/utils-1.10.0.tgz", + "integrity": "sha512-KLs73wIU/T3ZA+H+YTlzQ1fT+6p02RfixMQ7+l8S+IyLxc+nwSMABYnoZybJJuRvw1huJuFv1+n0B0HDteDFZA==", + "license": "Apache-2.0", + "dependencies": { + "@deephaven/log": "^1.8.0", + "nanoid": "^5.0.7" + }, + "engines": { + "node": ">=16" + } + }, "node_modules/@deephaven/log": { "version": "0.97.0", "resolved": "https://registry.npmjs.org/@deephaven/log/-/log-0.97.0.tgz", @@ -1532,6 +1594,16 @@ "node": ">=14" } }, + "node_modules/@react-types/shared": { + "version": "3.33.1", + "resolved": "https://registry.npmjs.org/@react-types/shared/-/shared-3.33.1.tgz", + "integrity": "sha512-oJHtjvLG43VjwemQDadlR5g/8VepK56B/xKO2XORPHt9zlW6IZs3tZrYlvH29BMvoqC7RtE7E5UjgbnbFtDGag==", + "dev": true, + "license": "Apache-2.0", + "peerDependencies": { + "react": "^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0-rc.1" + } + }, "node_modules/@redhat-developer/locators": { "version": "1.18.1", "resolved": "https://registry.npmjs.org/@redhat-developer/locators/-/locators-1.18.1.tgz", @@ -2392,6 +2464,16 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/react": { + "version": "19.2.14", + "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.14.tgz", + "integrity": "sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==", + "dev": true, + "license": "MIT", + "dependencies": { + "csstype": "^3.2.2" + } + }, "node_modules/@types/readdir-glob": { "version": "1.1.5", "resolved": "https://registry.npmjs.org/@types/readdir-glob/-/readdir-glob-1.1.5.tgz", @@ -5154,6 +5236,13 @@ "node": ">=18" } }, + "node_modules/csstype": { + "version": "3.2.3", + "resolved": "https://registry.npmjs.org/csstype/-/csstype-3.2.3.tgz", + "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", + "dev": true, + "license": "MIT" + }, "node_modules/ctrf": { "version": "0.0.9", "resolved": "https://registry.npmjs.org/ctrf/-/ctrf-0.0.9.tgz", @@ -8744,6 +8833,12 @@ "integrity": "sha512-LgVTMpQtIopCi79SJeDiP0TfWi5CNEc/L/aRdTh3yIvmZXTnheWpKjSZhnvMl8iXbC1tFg9gdHHDMLoV7CnG+w==", "license": "MIT" }, + "node_modules/lodash.clamp": { + "version": "4.0.3", + "resolved": "https://registry.npmjs.org/lodash.clamp/-/lodash.clamp-4.0.3.tgz", + "integrity": "sha512-HvzRFWjtcguTW7yd8NJBshuNaCa8aqNFtnswdT7f/cMd/1YKy5Zzoq4W/Oxvnx9l7aeY258uSdDfM793+eLsVg==", + "license": "MIT" + }, "node_modules/lodash.includes": { "version": "4.3.0", "resolved": "https://registry.npmjs.org/lodash.includes/-/lodash.includes-4.3.0.tgz", @@ -10768,6 +10863,17 @@ "node": ">=0.10.0" } }, + "node_modules/react": { + "version": "19.2.4", + "resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz", + "integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==", + "dev": true, + "license": "MIT", + "peer": true, + "engines": { + "node": ">=0.10.0" + } + }, "node_modules/read": { "version": "1.0.7", "resolved": "https://registry.npmjs.org/read/-/read-1.0.7.tgz", @@ -14910,6 +15016,11 @@ } } }, + "@deephaven/filters": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@deephaven/filters/-/filters-1.1.0.tgz", + "integrity": "sha512-SMWKYT8aFtZ/CyVUtUiF1/RPQGl4Y6dvNG43KvmIMKKMXYqhliKD291aynYSl7C8IafkHNzdDtdAZPgPXLoOWA==" + }, "@deephaven/jsapi-nodejs": { "version": "1.15.0", "resolved": "https://registry.npmjs.org/@deephaven/jsapi-nodejs/-/jsapi-nodejs-1.15.0.tgz", @@ -14953,6 +15064,45 @@ "integrity": "sha512-kz+CAwFctjmt1EwHrEfHt21kYGIGrjSQM3g4iSM04am3uoRMdbg6dO+cR6pxq+O+bNaC0/cJ6/yyn/FlXh/1xg==", "dev": true }, + "@deephaven/jsapi-utils": { + "version": "1.16.0", + "resolved": "https://registry.npmjs.org/@deephaven/jsapi-utils/-/jsapi-utils-1.16.0.tgz", + "integrity": "sha512-lyMDgxmb7v2GDlm1Ksf4S037QuWUhpOMKtnHZoEuKKIthcd6WejsNP+S8U2iRCYSJiJ1jEpTAVMrrS2CM4ZC7w==", + "requires": { + "@deephaven/filters": "^1.1.0", + "@deephaven/jsapi-types": "^1.0.0-dev0.40.4", + "@deephaven/log": "^1.8.0", + "@deephaven/utils": "^1.10.0", + "lodash.clamp": "^4.0.3", + "nanoid": "^5.0.7" + }, + "dependencies": { + "@deephaven/jsapi-types": { + "version": "1.0.0-dev0.40.9", + "resolved": "https://registry.npmjs.org/@deephaven/jsapi-types/-/jsapi-types-1.0.0-dev0.40.9.tgz", + "integrity": "sha512-NwMxFmNCnRV4/2+MdN/8vUGiEtXFgL1K/+iXTKKvi+Brje5JHOSCn2miCKR9tAn0LNb/UdmJq+DSIZqvz8cU/Q==" + }, + "@deephaven/log": { + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@deephaven/log/-/log-1.8.0.tgz", + "integrity": "sha512-gzp6/7qW4W8Re+DLSaG33KEQJ30OrrNq3cNQA8fUeXQrabSNOIsyeVOaerQ/57c4zWhWVKamplax0LIYRsmDiw==", + "requires": { + "event-target-shim": "^6.0.2", + "jszip": "^3.10.1", + "safe-stable-stringify": "^2.5.0" + } + }, + "@deephaven/utils": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@deephaven/utils/-/utils-1.10.0.tgz", + "integrity": "sha512-KLs73wIU/T3ZA+H+YTlzQ1fT+6p02RfixMQ7+l8S+IyLxc+nwSMABYnoZybJJuRvw1huJuFv1+n0B0HDteDFZA==", + "requires": { + "@deephaven/log": "^1.8.0", + "nanoid": "^5.0.7" + } + } + } + }, "@deephaven/log": { "version": "0.97.0", "resolved": "https://registry.npmjs.org/@deephaven/log/-/log-0.97.0.tgz", @@ -15483,6 +15633,13 @@ "integrity": "sha512-+1VkjdD0QBLPodGrJUeqarH8VAIvQODIbwh9XpP5Syisf7YoQgsJKPNFoqqLQlu+VQ/tVSshMR6loPMn8U+dPg==", "optional": true }, + "@react-types/shared": { + "version": "3.33.1", + "resolved": "https://registry.npmjs.org/@react-types/shared/-/shared-3.33.1.tgz", + "integrity": "sha512-oJHtjvLG43VjwemQDadlR5g/8VepK56B/xKO2XORPHt9zlW6IZs3tZrYlvH29BMvoqC7RtE7E5UjgbnbFtDGag==", + "dev": true, + "requires": {} + }, "@redhat-developer/locators": { "version": "1.18.1", "resolved": "https://registry.npmjs.org/@redhat-developer/locators/-/locators-1.18.1.tgz", @@ -16063,6 +16220,15 @@ "integrity": "sha512-37i+OaWTh9qeK4LSHPsyRC7NahnGotNuZvjLSgcPzblpHB3rrCJxAOgI5gCdKm7coonsaX1Of0ILiTcnZjbfxA==", "dev": true }, + "@types/react": { + "version": "19.2.14", + "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.14.tgz", + "integrity": "sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==", + "dev": true, + "requires": { + "csstype": "^3.2.2" + } + }, "@types/readdir-glob": { "version": "1.1.5", "resolved": "https://registry.npmjs.org/@types/readdir-glob/-/readdir-glob-1.1.5.tgz", @@ -17915,6 +18081,12 @@ "rrweb-cssom": "^0.8.0" } }, + "csstype": { + "version": "3.2.3", + "resolved": "https://registry.npmjs.org/csstype/-/csstype-3.2.3.tgz", + "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", + "dev": true + }, "ctrf": { "version": "0.0.9", "resolved": "https://registry.npmjs.org/ctrf/-/ctrf-0.0.9.tgz", @@ -20455,6 +20627,11 @@ "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.23.tgz", "integrity": "sha512-LgVTMpQtIopCi79SJeDiP0TfWi5CNEc/L/aRdTh3yIvmZXTnheWpKjSZhnvMl8iXbC1tFg9gdHHDMLoV7CnG+w==" }, + "lodash.clamp": { + "version": "4.0.3", + "resolved": "https://registry.npmjs.org/lodash.clamp/-/lodash.clamp-4.0.3.tgz", + "integrity": "sha512-HvzRFWjtcguTW7yd8NJBshuNaCa8aqNFtnswdT7f/cMd/1YKy5Zzoq4W/Oxvnx9l7aeY258uSdDfM793+eLsVg==" + }, "lodash.includes": { "version": "4.3.0", "resolved": "https://registry.npmjs.org/lodash.includes/-/lodash.includes-4.3.0.tgz", @@ -21881,6 +22058,13 @@ } } }, + "react": { + "version": "19.2.4", + "resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz", + "integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==", + "dev": true, + "peer": true + }, "read": { "version": "1.0.7", "resolved": "https://registry.npmjs.org/read/-/read-1.0.7.tgz", diff --git a/package.json b/package.json index 59286f537..5e678727b 100644 --- a/package.json +++ b/package.json @@ -1066,6 +1066,7 @@ "@deephaven-enterprise/auth-nodejs": "^1.20250219.134-beta", "@deephaven-enterprise/query-utils": "^1.20250219.134-beta", "@deephaven/jsapi-nodejs": "^1.15.0", + "@deephaven/jsapi-utils": "^1.16.0", "@modelcontextprotocol/sdk": "^1.27.1", "archiver": "^7.0.1", "chai": "^4.5.0", @@ -1076,10 +1077,12 @@ "@cfworker/json-schema": "^4.1.1", "@deephaven-enterprise/jsapi-types": "^1.20250219.134-beta", "@deephaven/jsapi-types": "^41.2.0", + "@react-types/shared": "^3.33.1", "@types/archiver": "^6.0.3", "@types/chai": "^4.3.20", "@types/mocha": "^10.0.10", "@types/node": "22.19.1", + "@types/react": "^19.2.14", "@types/vscode": "^1.105.0", "@types/vscode-webview": "^1.57.5", "@types/ws": "^8.5.10", diff --git a/src/mcp/McpServer.spec.ts b/src/mcp/McpServer.spec.ts new file mode 100644 index 000000000..2a9cfdea1 --- /dev/null +++ b/src/mcp/McpServer.spec.ts @@ -0,0 +1,439 @@ +import * as http from 'http'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { McpTool, McpToolSpec } from '../types'; +import type { OutputChannelWithHistory } from '../util'; + +/* eslint-disable @typescript-eslint/naming-convention */ +// HTTP headers and MCP protocol headers must use their spec-defined names + +vi.mock('vscode'); + +// Mock all tool creators to return minimal stub tools +vi.mock('./tools', () => ({ + createAddRemoteFileSourcesTool: vi.fn(() => + makeStubTool('addRemoteFileSources') + ), + createConnectToServerTool: vi.fn(() => makeStubTool('connectToServer')), + createGetColumnStatsTool: vi.fn(() => makeStubTool('getColumnStats')), + createGetLogsTool: vi.fn(() => makeStubTool('getLogs')), + createGetTableDataTool: vi.fn(() => makeStubTool('getTableData')), + createGetTableStatsTool: vi.fn(() => makeStubTool('getTableStats')), + createListConnectionsTool: vi.fn(() => makeStubTool('listConnections')), + createListVariablesTool: vi.fn(() => makeStubTool('listVariables')), + createListRemoteFileSourcesTool: vi.fn(() => + makeStubTool('listRemoteFileSources') + ), + createListServersTool: vi.fn(() => makeStubTool('listServers')), + createOpenFilesInEditorTool: vi.fn(() => makeStubTool('openFilesInEditor')), + createOpenVariablePanelsTool: vi.fn(() => makeStubTool('openVariablePanels')), + createRemoveRemoteFileSourcesTool: vi.fn(() => + makeStubTool('removeRemoteFileSources') + ), + createRunCodeFromUriTool: vi.fn(() => makeStubTool('runCodeFromUri')), + createRunCodeTool: vi.fn(() => makeStubTool('runCode')), + createSetEditorConnectionTool: vi.fn(() => + makeStubTool('setEditorConnection') + ), + createShowOutputPanelTool: vi.fn(() => makeStubTool('showOutputPanel')), +})); + +vi.mock('./tools/connectToServer', () => ({ + createConnectToServerTool: vi.fn(() => makeStubTool('connectToServer')), +})); + +/** Create a minimal stub tool that echoes back its args */ +function makeStubTool(name: string): McpTool { + return { + name, + spec: { + title: name, + description: `Stub tool: ${name}`, + inputSchema: {}, + outputSchema: {}, + }, + handler: vi.fn().mockResolvedValue({ + content: [{ type: 'text', text: `stub result for ${name}` }], + }), + }; +} + +/** Create a minimal mock for OutputChannelWithHistory */ +function makeMockOutputChannel(): OutputChannelWithHistory { + return { + appendLine: vi.fn(), + append: vi.fn(), + show: vi.fn(), + clear: vi.fn(), + dispose: vi.fn(), + hide: vi.fn(), + replace: vi.fn(), + name: 'mock', + } as unknown as OutputChannelWithHistory; +} + +// Helper: send an HTTP POST to the MCP server and collect the full response +function postToMcp( + port: number, + body: unknown, + headers: Record = {} +): Promise<{ + status: number; + headers: http.IncomingMessage['headers']; + body: string; +}> { + return new Promise((resolve, reject) => { + const payload = JSON.stringify(body); + const req = http.request( + { + hostname: '127.0.0.1', + port, + path: '/mcp', + method: 'POST', + headers: { + 'Content-Type': 'application/json', + // StreamableHTTPServerTransport checks for acceptable response formats. + // Include both JSON and SSE to satisfy the transport's Accept check. + Accept: 'application/json, text/event-stream', + 'Content-Length': Buffer.byteLength(payload), + ...headers, + }, + }, + res => { + let data = ''; + res.on('data', chunk => (data += chunk)); + res.on('end', () => + resolve({ + status: res.statusCode ?? 0, + headers: res.headers, + body: data, + }) + ); + } + ); + req.on('error', reject); + req.write(payload); + req.end(); + }); +} + +const INITIALIZE_REQUEST = { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: '2024-11-05', + capabilities: {}, + clientInfo: { name: 'test-client', version: '0.0.1' }, + }, +}; + +const LIST_TOOLS_REQUEST = { + jsonrpc: '2.0', + id: 2, + method: 'tools/list', + params: {}, +}; + +describe('McpServer', () => { + // Dynamic import so mocks above apply + let McpServer: (typeof import('./McpServer'))['McpServer']; + let server: import('./McpServer').McpServer; + let port: number; + + beforeEach(async () => { + ({ McpServer } = await import('./McpServer')); + + server = new McpServer( + {} as any, // coreJsApiCache + makeMockOutputChannel() as any, // outputChannel + makeMockOutputChannel() as any, // outputChannelDebug + {} as any, // panelService + {} as any, // pythonDiagnostics + {} as any, // pythonWorkspace + {} as any // serverManager + ); + + port = await server.start(); + }); + + afterEach(async () => { + await server.stop(); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 1. Initialize request creates a new session + // ───────────────────────────────────────────────────────────────────────── + describe('session initialization', () => { + it('should accept an initialize request and return a session ID', async () => { + const res = await postToMcp(port, INITIALIZE_REQUEST); + + expect(res.status).toBe(200); + // The session ID is returned either in the response header or body + const sessionId = + (res.headers['mcp-session-id'] as string | undefined) ?? + ((): string | undefined => { + try { + return JSON.parse(res.body)?.sessionId as string | undefined; + } catch { + return undefined; + } + })(); + + expect(sessionId).toBeDefined(); + expect(typeof sessionId).toBe('string'); + expect(sessionId!.length).toBeGreaterThan(0); + }); + + it('should reject a non-initialize request with no session ID', async () => { + const res = await postToMcp(port, LIST_TOOLS_REQUEST); + + expect(res.status).toBe(400); + const body = JSON.parse(res.body); + expect(body.error).toBeDefined(); + expect(body.error.code).toBe(-32600); + }); + + it('should reject unsupported methods (PUT, DELETE, etc.) with 405', async () => { + const res = await new Promise<{ status: number }>((resolve, reject) => { + const req = http.request( + { hostname: '127.0.0.1', port, path: '/mcp', method: 'PUT' }, + res => resolve({ status: res.statusCode ?? 0 }) + ); + req.on('error', reject); + req.end(); + }); + + expect(res.status).toBe(405); + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 2. Session persistence — same session ID reuses the session + // ───────────────────────────────────────────────────────────────────────── + describe('session persistence', () => { + it('should reuse an existing session for subsequent requests', async () => { + // Step 1: initialize + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + expect(initRes.status).toBe(200); + const sessionId = initRes.headers['mcp-session-id'] as string; + expect(sessionId).toBeDefined(); + + // Step 2: send tools/list with the same session ID + const toolsRes = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': sessionId, + }); + + expect(toolsRes.status).toBe(200); + const toolsBody = JSON.parse(toolsRes.body); + // Should have a result (not an error) + expect(toolsBody.error).toBeUndefined(); + expect(toolsBody.result).toBeDefined(); + }); + + it('should return all registered tools via tools/list', async () => { + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + const sessionId = initRes.headers['mcp-session-id'] as string; + + const toolsRes = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': sessionId, + }); + + const toolsBody = JSON.parse(toolsRes.body); + expect(toolsBody.result?.tools).toBeInstanceOf(Array); + // All 17 stub tools should be present + expect(toolsBody.result.tools.length).toBe(17); + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 3. Session isolation — different session IDs get isolated server instances + // ───────────────────────────────────────────────────────────────────────── + describe('session isolation', () => { + it('should create independent sessions for separate initialize requests', async () => { + const [res1, res2] = await Promise.all([ + postToMcp(port, INITIALIZE_REQUEST), + postToMcp(port, INITIALIZE_REQUEST), + ]); + + expect(res1.status).toBe(200); + expect(res2.status).toBe(200); + + const session1 = res1.headers['mcp-session-id'] as string; + const session2 = res2.headers['mcp-session-id'] as string; + + expect(session1).toBeDefined(); + expect(session2).toBeDefined(); + // Sessions should be distinct + expect(session1).not.toBe(session2); + }); + + it('should reject requests using an unknown session ID', async () => { + const res = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': 'non-existent-session-id', + }); + + // Either 400 (bad session) or the server sends back an error JSON + expect([400, 200]).toContain(res.status); + if (res.status === 200) { + const body = JSON.parse(res.body); + // If 200, should contain an error in the JSON-RPC response + expect(body.error).toBeDefined(); + } + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 4. Parallel requests — no race conditions + // ───────────────────────────────────────────────────────────────────────── + describe('parallel requests', () => { + it('should handle parallel tool calls on the same session without errors', async () => { + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + const sessionId = initRes.headers['mcp-session-id'] as string; + + // Fire 5 concurrent tools/list requests on the same session + const results = await Promise.all( + Array.from({ length: 5 }, (_, i) => + postToMcp( + port, + { ...LIST_TOOLS_REQUEST, id: i + 10 }, + { 'mcp-session-id': sessionId } + ) + ) + ); + + // All should succeed without server errors + for (const res of results) { + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.error).toBeUndefined(); + expect(body.result).toBeDefined(); + } + }); + + it('should handle parallel requests across different sessions without errors', async () => { + // Initialize two sessions simultaneously + const [init1, init2] = await Promise.all([ + postToMcp(port, INITIALIZE_REQUEST), + postToMcp(port, INITIALIZE_REQUEST), + ]); + + const session1 = init1.headers['mcp-session-id'] as string; + const session2 = init2.headers['mcp-session-id'] as string; + + // Then fire requests on both sessions in parallel + const [res1a, res1b, res2a, res2b] = await Promise.all([ + postToMcp( + port, + { ...LIST_TOOLS_REQUEST, id: 21 }, + { 'mcp-session-id': session1 } + ), + postToMcp( + port, + { ...LIST_TOOLS_REQUEST, id: 22 }, + { 'mcp-session-id': session1 } + ), + postToMcp( + port, + { ...LIST_TOOLS_REQUEST, id: 23 }, + { 'mcp-session-id': session2 } + ), + postToMcp( + port, + { ...LIST_TOOLS_REQUEST, id: 24 }, + { 'mcp-session-id': session2 } + ), + ]); + + for (const res of [res1a, res1b, res2a, res2b]) { + expect(res.status).toBe(200); + const body = JSON.parse(res.body); + expect(body.error).toBeUndefined(); + expect(body.result).toBeDefined(); + } + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 5. Tools functionality + // ───────────────────────────────────────────────────────────────────────── + describe('MCP tools availability', () => { + it('should have addRemoteFileSources tool registered', async () => { + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + const sessionId = initRes.headers['mcp-session-id'] as string; + + const toolsRes = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': sessionId, + }); + const tools: { name: string }[] = + JSON.parse(toolsRes.body).result?.tools ?? []; + expect(tools.some(t => t.name === 'addRemoteFileSources')).toBe(true); + }); + + it('should have listServers tool registered', async () => { + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + const sessionId = initRes.headers['mcp-session-id'] as string; + + const toolsRes = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': sessionId, + }); + const tools: { name: string }[] = + JSON.parse(toolsRes.body).result?.tools ?? []; + expect(tools.some(t => t.name === 'listServers')).toBe(true); + }); + + it('should have runCode tool registered', async () => { + const initRes = await postToMcp(port, INITIALIZE_REQUEST); + const sessionId = initRes.headers['mcp-session-id'] as string; + + const toolsRes = await postToMcp(port, LIST_TOOLS_REQUEST, { + 'mcp-session-id': sessionId, + }); + const tools: { name: string }[] = + JSON.parse(toolsRes.body).result?.tools ?? []; + expect(tools.some(t => t.name === 'runCode')).toBe(true); + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 6. Session cleanup + // ───────────────────────────────────────────────────────────────────────── + describe('session cleanup on stop()', () => { + it('should clear all sessions when stop() is called', async () => { + // Create two sessions + await Promise.all([ + postToMcp(port, INITIALIZE_REQUEST), + postToMcp(port, INITIALIZE_REQUEST), + ]); + + // Access internal Maps for verification + const mcpServer = server as unknown as { + transports: Map; + servers: Map; + }; + + expect(mcpServer.transports.size).toBe(2); + expect(mcpServer.servers.size).toBe(2); + + await server.stop(); + + expect(mcpServer.transports.size).toBe(0); + expect(mcpServer.servers.size).toBe(0); + }); + }); + + // ───────────────────────────────────────────────────────────────────────── + // 7. Port allocation + // ───────────────────────────────────────────────────────────────────────── + describe('port management', () => { + it('should return the allocated port via getPort()', () => { + expect(server.getPort()).toBe(port); + expect(typeof port).toBe('number'); + expect(port).toBeGreaterThan(0); + }); + + it('should return null from getPort() after stop()', async () => { + await server.stop(); + expect(server.getPort()).toBeNull(); + }); + }); +}); diff --git a/src/mcp/McpServer.ts b/src/mcp/McpServer.ts index 56fd3dd70..6b90315e5 100644 --- a/src/mcp/McpServer.ts +++ b/src/mcp/McpServer.ts @@ -1,7 +1,9 @@ import * as vscode from 'vscode'; +import { randomUUID } from 'node:crypto'; import type { dh as DhcType } from '@deephaven/jsapi-types'; import { McpServer as SdkMcpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'; import * as http from 'http'; import type { IAsyncCacheService, @@ -38,9 +40,10 @@ import { createConnectToServerTool } from './tools/connectToServer'; * Provides tools for AI assistants (like GitHub Copilot) to interact with Deephaven. */ export class McpServer extends DisposableBase { - private server: SdkMcpServer; private httpServer: http.Server | null = null; private port: number | null = null; + private transports: Map = new Map(); + private servers: Map = new Map(); constructor( readonly coreJsApiCache: IAsyncCacheService, @@ -52,43 +55,162 @@ export class McpServer extends DisposableBase { readonly serverManager: IServerManager ) { super(); + } - // Create an MCP server - this.server = new SdkMcpServer({ + private createServer(): SdkMcpServer { + const server = new SdkMcpServer({ name: MCP_SERVER_NAME, version: '1.0.0', }); - this.registerTool(createAddRemoteFileSourcesTool()); - this.registerTool(createConnectToServerTool(this)); - this.registerTool(createGetColumnStatsTool(this)); - this.registerTool(createGetLogsTool(this)); - this.registerTool(createGetTableDataTool(this)); - this.registerTool(createGetTableStatsTool(this)); - this.registerTool(createListConnectionsTool(this)); - this.registerTool(createListVariablesTool(this)); - this.registerTool(createListRemoteFileSourcesTool(this)); - this.registerTool(createListServersTool(this)); - this.registerTool(createOpenFilesInEditorTool()); - this.registerTool(createOpenVariablePanelsTool(this)); - this.registerTool(createRemoveRemoteFileSourcesTool()); - this.registerTool(createRunCodeFromUriTool(this)); - this.registerTool(createRunCodeTool(this)); - this.registerTool(createSetEditorConnectionTool(this)); - this.registerTool(createShowOutputPanelTool(this)); + this.registerToolsOnServer(server); + + return server; + } + + private registerToolsOnServer(server: SdkMcpServer): void { + this.registerTool(server, createAddRemoteFileSourcesTool()); + this.registerTool(server, createConnectToServerTool(this)); + this.registerTool(server, createGetColumnStatsTool(this)); + this.registerTool(server, createGetLogsTool(this)); + this.registerTool(server, createGetTableDataTool(this)); + this.registerTool(server, createGetTableStatsTool(this)); + this.registerTool(server, createListConnectionsTool(this)); + this.registerTool(server, createListVariablesTool(this)); + this.registerTool(server, createListRemoteFileSourcesTool(this)); + this.registerTool(server, createListServersTool(this)); + this.registerTool(server, createOpenFilesInEditorTool()); + this.registerTool(server, createOpenVariablePanelsTool(this)); + this.registerTool(server, createRemoveRemoteFileSourcesTool()); + this.registerTool(server, createRunCodeFromUriTool(this)); + this.registerTool(server, createRunCodeTool(this)); + this.registerTool(server, createSetEditorConnectionTool(this)); + this.registerTool(server, createShowOutputPanelTool(this)); + } + + private registerTool( + server: SdkMcpServer, + { name, spec, handler }: McpTool + ): void { + server.registerTool(name, spec, handler); + } + + private handleInvalidPath(res: http.ServerResponse): void { + res.writeHead(404, { contentType: 'text/plain' }); + res.end('Not found'); + } + + private handleInvalidMethod(res: http.ServerResponse): void { + res.writeHead(405, { + contentType: 'text/plain', + allow: 'GET, POST', + }); + res.end('Method Not Allowed'); + } + + private handleSessionNotFound( + res: http.ServerResponse, + sessionId: string + ): void { + // eslint-disable-next-line @typescript-eslint/naming-convention + res.writeHead(400, { 'Content-Type': 'application/json' }); + res.end( + JSON.stringify({ + jsonrpc: '2.0', + error: { + code: -32600, + message: `Bad Request: session ID ${sessionId} not found`, + }, + id: null, + }) + ); + } + + private async handleExistingSession( + req: http.IncomingMessage, + res: http.ServerResponse, + requestBody: unknown, + sessionId: string + ): Promise { + if (!this.transports.has(sessionId)) { + this.handleSessionNotFound(res, sessionId); + return; + } + + // Existing session — reuse transport + const transport = this.transports.get(sessionId)!; + await transport.handleRequest(req, res, requestBody); } - private registerTool({ - name, - spec, - handler, - }: McpTool): void { - this.server.registerTool(name, spec, handler); + private handleInvalidSessionIdForRequestType( + res: http.ServerResponse, + message: string + ): void { + // eslint-disable-next-line @typescript-eslint/naming-convention + res.writeHead(400, { 'Content-Type': 'application/json' }); + res.end( + JSON.stringify({ + jsonrpc: '2.0', + error: { + code: -32600, + message, + }, + id: null, + }) + ); + } + + private async handleNewSession( + req: http.IncomingMessage, + res: http.ServerResponse, + requestBody: unknown + ): Promise { + // New session — create isolated server + transport pair + const server = this.createServer(); + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: (): string => randomUUID(), + enableJsonResponse: true, + onsessioninitialized: (sid): void => { + this.transports.set(sid, transport); + this.servers.set(sid, server); + }, + }); + + transport.onclose = async (): Promise => { + try { + const sid = transport.sessionId; + if (sid) { + this.transports.delete(sid); + const closingServer = this.servers.get(sid); + this.servers.delete(sid); + await closingServer?.close(); + } + } catch (error) { + this.outputChannelDebug.appendLine( + `[McpServer] Error during session cleanup: ${error instanceof Error ? error.message : String(error)}` + ); + } + }; + + await server.connect(transport); + await transport.handleRequest(req, res, requestBody); + } + + private handleRequestError(res: http.ServerResponse, error: unknown): void { + res.writeHead(500, { contentType: 'application/json' }); + res.end( + JSON.stringify({ + error: `Failed to process request: ${error instanceof Error ? error.message : String(error)}`, + }) + ); } /** * Start the MCP server on an HTTP endpoint. - * Creates a new transport for each request (stateless operation). + * Uses stateful session management: each initialize request creates a new + * isolated server+transport pair stored by session ID. Subsequent requests + * from the same session reuse the existing transport, eliminating race + * conditions from multiple concurrent requests. * * @param preferredPort Optional port to try first. If not provided or unavailable, will auto-allocate. * @returns The actual port the server is listening on @@ -100,50 +222,59 @@ export class McpServer extends DisposableBase { this.httpServer = http.createServer(async (req, res) => { if (req.url !== '/mcp') { - res.writeHead(404, { contentType: 'text/plain' }); - res.end('Not found'); + this.handleInvalidPath(res); + return; } - // Only accept POST requests since we don't currenlty support SSE. TBD - // whether we need SSE in the future. - if (req.method !== 'POST') { - res.writeHead(405, { - contentType: 'text/plain', - allow: 'POST', - }); - res.end('Method Not Allowed'); + // Accept GET (for SSE) and POST (for JSON-RPC) requests. + // Other methods are not supported by the MCP protocol. + if (req.method !== 'GET' && req.method !== 'POST') { + this.handleInvalidMethod(res); return; } - // Collect the request body + // Collect body only for POST requests let body = ''; - req.on('data', chunk => { - body += chunk.toString(); - }); + if (req.method === 'POST') { + req.on('data', chunk => { + body += chunk.toString(); + }); + } req.on('end', async () => { try { - const requestBody = JSON.parse(body); + // Parse body if present, otherwise undefined for GET requests + const requestBody = body ? JSON.parse(body) : undefined; + const sessionId = req.headers['mcp-session-id'] as string | undefined; + const hasSessionId = sessionId != null; - // Create a new transport for each request - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: undefined, - enableJsonResponse: true, - }); + // Only POST requests can be an initialize request + const isInitializeReq = requestBody + ? isInitializeRequest(requestBody) + : false; - res.on('close', () => { - transport.close(); - }); + // Validate: initialize requests must NOT have a session ID, + // and non-initialize requests MUST have a session ID. + if (hasSessionId === isInitializeReq) { + this.handleInvalidSessionIdForRequestType( + res, + hasSessionId + ? 'Bad Request: initialize request must not include mcp-session-id header' + : 'Bad Request: include mcp-session-id header for existing sessions, or send an initialize request to start a new session' + ); + return; + } - await this.server.connect(transport); - await transport.handleRequest(req, res, requestBody); + sessionId == null + ? await this.handleNewSession(req, res, requestBody) + : await this.handleExistingSession( + req, + res, + requestBody, + sessionId + ); } catch (error) { - res.writeHead(500, { contentType: 'application/json' }); - res.end( - JSON.stringify({ - error: `Failed to process request: ${error instanceof Error ? error.message : String(error)}`, - }) - ); + this.handleRequestError(res, error); } }); }); @@ -207,6 +338,20 @@ export class McpServer extends DisposableBase { return; } + // Close all active sessions before shutting down the HTTP server + for (const [sid, transport] of this.transports) { + try { + await transport.close(); + await this.servers.get(sid)?.close(); + } catch (error) { + this.outputChannelDebug.appendLine( + `[McpServer] Error closing session ${sid}: ${error instanceof Error ? error.message : String(error)}` + ); + } + } + this.transports.clear(); + this.servers.clear(); + const { resolve, reject, promise } = withResolvers(); this.httpServer.close(err => { diff --git a/src/mcp/tools/connectToServer.ts b/src/mcp/tools/connectToServer.ts index ce2db8a19..a56f95795 100644 --- a/src/mcp/tools/connectToServer.ts +++ b/src/mcp/tools/connectToServer.ts @@ -16,7 +16,10 @@ const spec = { inputSchema: { url: z.string().describe('Server URL (e.g., "http://localhost:10000")'), }, - outputSchema: createMcpToolOutputSchema(), + outputSchema: createMcpToolOutputSchema({ + type: z.enum(['DHC', 'DHE']).optional().describe('Server type'), + url: z.string().optional().describe('Server URL'), + }), } as const; type Spec = typeof spec; diff --git a/src/mcp/tools/getColumnStats.spec.ts b/src/mcp/tools/getColumnStats.spec.ts index 0b1a57e11..0b865e8d5 100644 --- a/src/mcp/tools/getColumnStats.spec.ts +++ b/src/mcp/tools/getColumnStats.spec.ts @@ -2,17 +2,22 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { dh as DhcType } from '@deephaven/jsapi-types'; import { createGetColumnStatsTool } from './getColumnStats'; -import type { IServerManager, ServerState } from '../../types'; -import { DhcService } from '../../services'; +import type { IAsyncCacheService, IServerManager } from '../../types'; import { fakeMcpToolTimings, mcpSuccessResult, - mcpErrorResult, MOCK_DHC_URL, } from '../utils/mcpTestUtils'; +import { getTableOrError } from '../utils/tableUtils'; vi.mock('vscode'); -vi.mock('../../services/DhcService'); +vi.mock('../utils/tableUtils', () => ({ + getTableOrError: vi.fn(), + convertColumnStatsToRecords: vi.fn(stats => ({ + statistics: Object.fromEntries(stats.statisticsMap), + uniqueValues: Object.fromEntries(stats.uniqueValues), + })), +})); const MOCK_COLUMN = { name: 'Price', @@ -85,23 +90,14 @@ const EXPECTED_SUCCESS_NO_UNIQUE = mcpSuccessResult('Column stats retrieved', { }); /* eslint-enable @typescript-eslint/naming-convention */ -const MOCK_SERVER_RUNNING: ServerState = { - isRunning: true, - type: 'DHC', - url: MOCK_DHC_URL, - isConnected: false, - connectionCount: 0, -}; - describe('createGetColumnStatsTool', () => { - const mockSession = { - getObject: vi.fn(), - } as unknown as DhcType.IdeSession; - - const mockConnection = Object.assign(Object.create(DhcService.prototype), { - initSession: vi.fn(), - getSession: vi.fn(), - }) as DhcService; + const coreJsApiCache = { + get: vi.fn(), + has: vi.fn(), + invalidate: vi.fn(), + dispose: vi.fn(), + onDidInvalidate: vi.fn(), + } as unknown as IAsyncCacheService; const serverManager = { getServer: vi.fn(), @@ -113,18 +109,12 @@ describe('createGetColumnStatsTool', () => { vi.clearAllMocks(); fakeMcpToolTimings(); - // Default successful getServer mock - vi.mocked(serverManager.getServer).mockReturnValue(MOCK_SERVER_RUNNING); - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - // Mock isInitialized getter - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => true), - configurable: true, + // Default successful mocks + vi.mocked(getTableOrError).mockResolvedValue({ + success: true, + table: MOCK_TABLE, + connectionUrl: MOCK_DHC_URL, }); - - vi.mocked(mockConnection.getSession).mockResolvedValue(mockSession); - vi.mocked(mockSession.getObject).mockResolvedValue(MOCK_TABLE); vi.mocked(MOCK_TABLE.findColumn).mockReturnValue(MOCK_COLUMN); vi.mocked(MOCK_TABLE.getColumnStatistics).mockResolvedValue( MOCK_COLUMN_STATS @@ -132,7 +122,7 @@ describe('createGetColumnStatsTool', () => { }); it('should return correct tool spec', () => { - const tool = createGetColumnStatsTool({ serverManager }); + const tool = createGetColumnStatsTool({ coreJsApiCache, serverManager }); expect(tool.name).toBe('getColumnStats'); expect(tool.spec.title).toBe('Get Column Statistics'); @@ -158,16 +148,22 @@ describe('createGetColumnStatsTool', () => { async ({ mockStats, expected }) => { vi.mocked(MOCK_TABLE.getColumnStatistics).mockResolvedValue(mockStats); - const tool = createGetColumnStatsTool({ serverManager }); + const tool = createGetColumnStatsTool({ + coreJsApiCache, + serverManager, + }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', columnName: 'Price', }); - expect(mockSession.getObject).toHaveBeenCalledWith({ - type: 'Table', - name: 'myTable', + expect(getTableOrError).toHaveBeenCalledWith({ + coreJsApiCache, + connectionUrlStr: MOCK_DHC_URL.href, + variableId: undefined, + tableName: 'myTable', + serverManager, }); expect(MOCK_TABLE.findColumn).toHaveBeenCalledWith('Price'); expect(MOCK_TABLE.getColumnStatistics).toHaveBeenCalledWith( @@ -180,116 +176,50 @@ describe('createGetColumnStatsTool', () => { }); describe('error handling', () => { - it('should initialize session if not initialized', async () => { - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => false), - configurable: true, - }); - - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - const tool = createGetColumnStatsTool({ serverManager }); - await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - columnName: 'Price', - }); - - expect(mockConnection.getSession).toHaveBeenCalled(); - }); - - it.each([ - { - name: 'invalid URL', - connectionUrl: 'invalid-url', - server: undefined, - connections: [], - expected: mcpErrorResult('Invalid URL: Invalid URL', { - connectionUrl: 'invalid-url', - tableName: 'myTable', - }), - shouldCallGetServer: false, - }, - { - name: 'missing connection', - connectionUrl: MOCK_DHC_URL.href, - server: undefined, - connections: [], - expected: mcpErrorResult('No connections or server found', { + it('should handle errors when getting table fails', async () => { + vi.mocked(getTableOrError).mockResolvedValue({ + success: false, + errorMessage: 'Connection failed', + details: { connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', - }), - shouldCallGetServer: true, - }, - { - name: 'missing session', - connectionUrl: MOCK_DHC_URL.href, - server: MOCK_SERVER_RUNNING, - connections: [mockConnection], - expected: mcpErrorResult('Unable to access session', { - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - }), - shouldCallGetServer: true, - }, - ])( - 'should handle $name', - async ({ - connectionUrl, - server, - connections, - expected, - shouldCallGetServer, - }) => { - vi.mocked(serverManager.getServer).mockReturnValue(server); - vi.mocked(serverManager.getConnections).mockReturnValue(connections); - vi.mocked(mockConnection.getSession).mockResolvedValue(null); - - const tool = createGetColumnStatsTool({ serverManager }); - const result = await tool.handler({ - connectionUrl, - tableName: 'myTable', - columnName: 'Price', - }); - - expect(result.structuredContent).toEqual(expected); - if (!shouldCallGetServer) { - expect(serverManager.getServer).not.toHaveBeenCalled(); - } - } - ); - - it('should handle errors and close table', async () => { - const error = new Error('Failed to get stats'); - vi.mocked(MOCK_TABLE.getColumnStatistics).mockRejectedValue(error); + }, + }); - const tool = createGetColumnStatsTool({ serverManager }); + const tool = createGetColumnStatsTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', columnName: 'Price', }); - expect(MOCK_TABLE.close).toHaveBeenCalled(); - expect(result.structuredContent).toEqual( - mcpErrorResult('Failed to get column stats: Failed to get stats', { + expect(result.structuredContent).toMatchObject({ + success: false, + message: 'Connection failed', + details: { connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', - columnName: 'Price', - }) - ); + }, + }); }); - it('should close table even on success', async () => { - const tool = createGetColumnStatsTool({ serverManager }); + it('should handle column not found error', async () => { + vi.mocked(MOCK_TABLE.findColumn).mockImplementation(() => { + throw new Error('Column not found'); + }); + + const tool = createGetColumnStatsTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', - columnName: 'Price', + columnName: 'InvalidColumn', }); - expect(result.structuredContent.success).toBe(true); expect(MOCK_TABLE.close).toHaveBeenCalled(); + expect(result.structuredContent).toMatchObject({ + success: false, + message: 'Failed to get column stats: Column not found', + }); }); }); }); diff --git a/src/mcp/tools/getColumnStats.ts b/src/mcp/tools/getColumnStats.ts index ef46c00f5..ff6fc588f 100644 --- a/src/mcp/tools/getColumnStats.ts +++ b/src/mcp/tools/getColumnStats.ts @@ -5,6 +5,7 @@ import type { McpToolHandlerArg, McpToolHandlerResult, IServerManager, + IAsyncCacheService, } from '../../types'; import { createMcpToolOutputSchema, @@ -63,8 +64,10 @@ type HandlerResult = McpToolHandlerResult; type GetColumnStatsTool = McpTool; export function createGetColumnStatsTool({ + coreJsApiCache, serverManager, }: { + coreJsApiCache: IAsyncCacheService; serverManager: IServerManager; }): GetColumnStatsTool { return { @@ -80,6 +83,7 @@ export function createGetColumnStatsTool({ try { const tableResult = await getTableOrError({ + coreJsApiCache, connectionUrlStr, variableId, tableName, diff --git a/src/mcp/tools/getTableData.spec.ts b/src/mcp/tools/getTableData.spec.ts index 78100dc9e..4c9878f80 100644 --- a/src/mcp/tools/getTableData.spec.ts +++ b/src/mcp/tools/getTableData.spec.ts @@ -1,45 +1,24 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { dh as DhcType } from '@deephaven/jsapi-types'; -import { createGetTableDataTool } from './getTableData'; -import type { IServerManager, ServerState } from '../../types'; -import { DhcService } from '../../services'; +import { + createGetTableDataTool, + DEFAULT_TABLE_PAGE_DATA_LIMIT, + DEFAULT_TABLE_PAGE_DATA_OFFSET, +} from './getTableData'; +import type { IAsyncCacheService, IServerManager } from '../../types'; import { fakeMcpToolTimings, mcpSuccessResult, - mcpErrorResult, MOCK_DHC_URL, } from '../utils/mcpTestUtils'; +import { getTableOrError, getTablePage } from '../utils/tableUtils'; vi.mock('vscode'); -vi.mock('../../services/DhcService'); - -const MOCK_VIEWPORT_DATA = { - rows: [ - { - get: vi.fn((col: DhcType.Column) => { - if (col.name === 'Symbol') { - return 'AAPL'; - } - if (col.name === 'Price') { - return 150.25; - } - return null; - }), - }, - { - get: vi.fn((col: DhcType.Column) => { - if (col.name === 'Symbol') { - return 'GOOGL'; - } - if (col.name === 'Price') { - return 2800.5; - } - return null; - }), - }, - ], -} as unknown as DhcType.ViewportData; +vi.mock('../utils/tableUtils', () => ({ + getTableOrError: vi.fn(), + getTablePage: vi.fn(), +})); const MOCK_TABLE = { size: 2, @@ -52,23 +31,41 @@ const MOCK_TABLE = { getViewportData: vi.fn(), } as unknown as DhcType.Table; -const MOCK_SERVER_RUNNING: ServerState = { - isRunning: true, - type: 'DHC', - url: MOCK_DHC_URL, - isConnected: false, - connectionCount: 0, +const MOCK_PAGE_DATA = { + columns: [ + { name: 'Symbol', type: 'java.lang.String' }, + { name: 'Price', type: 'double' }, + ], + data: [ + /* eslint-disable @typescript-eslint/naming-convention */ + { Symbol: 'AAPL', Price: 150.25 }, + { Symbol: 'GOOGL', Price: 2800.5 }, + /* eslint-enable @typescript-eslint/naming-convention */ + ], + hasMore: false, + rowCount: 2, + totalRows: 2, }; describe('getTableData', () => { - const mockSession = { - getObject: vi.fn(), - } as unknown as DhcType.IdeSession; + /* eslint-disable @typescript-eslint/naming-convention */ + const mockDh = { + VariableType: { + TABLE: 'Table', + TREETABLE: 'TreeTable', + HIERARCHICALTABLE: 'HierarchicalTable', + PARTITIONEDTABLE: 'PartitionedTable', + }, + } as unknown as typeof DhcType; + /* eslint-enable @typescript-eslint/naming-convention */ - const mockConnection = Object.assign(Object.create(DhcService.prototype), { - initSession: vi.fn(), - getSession: vi.fn(), - }) as DhcService; + const coreJsApiCache = { + get: vi.fn().mockResolvedValue(mockDh), + has: vi.fn(), + invalidate: vi.fn(), + dispose: vi.fn(), + onDidInvalidate: vi.fn(), + } as unknown as IAsyncCacheService; const serverManager = { getServer: vi.fn(), @@ -80,23 +77,17 @@ describe('getTableData', () => { vi.clearAllMocks(); fakeMcpToolTimings(); - // Default successful getServer mock - vi.mocked(serverManager.getServer).mockReturnValue(MOCK_SERVER_RUNNING); - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - // Mock isInitialized getter - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => true), - configurable: true, + // Default successful mocks + vi.mocked(getTableOrError).mockResolvedValue({ + success: true, + table: MOCK_TABLE, + connectionUrl: MOCK_DHC_URL, }); - - vi.mocked(mockConnection.getSession).mockResolvedValue(mockSession); - vi.mocked(mockSession.getObject).mockResolvedValue(MOCK_TABLE); - vi.mocked(MOCK_TABLE.getViewportData).mockResolvedValue(MOCK_VIEWPORT_DATA); + vi.mocked(getTablePage).mockResolvedValue(MOCK_PAGE_DATA); }); it('should return correct tool spec', () => { - const tool = createGetTableDataTool({ serverManager }); + const tool = createGetTableDataTool({ coreJsApiCache, serverManager }); expect(tool.name).toBe('getTableData'); expect(tool.spec.title).toBe('Get Table Data'); @@ -105,245 +96,128 @@ describe('getTableData', () => { ); }); - it('should successfully query table data with defaults', async () => { - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - }); - - expect(mockSession.getObject).toHaveBeenCalledWith({ - type: 'Table', - name: 'myTable', - }); - expect(MOCK_TABLE.setViewport).toHaveBeenCalledWith(0, 9); - expect(MOCK_TABLE.close).toHaveBeenCalled(); - expect(result.structuredContent).toEqual( - mcpSuccessResult('Fetched 2 rows', { - /* eslint-disable @typescript-eslint/naming-convention */ - columns: [ - { name: 'Symbol', type: 'java.lang.String' }, - { name: 'Price', type: 'double' }, - ], - connectionUrl: MOCK_DHC_URL.href, - data: [ - { Symbol: 'AAPL', Price: 150.25 }, - { Symbol: 'GOOGL', Price: 2800.5 }, - ], - hasMore: false, - limit: 10, - offset: 0, - rowCount: 2, - tableName: 'myTable', - totalRows: 2, - /* eslint-enable @typescript-eslint/naming-convention */ - }) - ); - }); - - it('should apply limit and show hasMore when more rows available', async () => { - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - - const largeTable = { + it.each([ + { + label: 'defaults', + tableSize: 2, + limit: DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset: DEFAULT_TABLE_PAGE_DATA_OFFSET, + hasMore: false, + }, + { + label: 'limit with hasMore when more rows available', + tableSize: 100, + limit: DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset: DEFAULT_TABLE_PAGE_DATA_OFFSET, + hasMore: true, + }, + { + label: 'custom limit', + tableSize: 100, + limit: 5, + offset: DEFAULT_TABLE_PAGE_DATA_OFFSET, + hasMore: true, + }, + { + label: 'offset for pagination', + tableSize: 100, + limit: DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset: 20, + hasMore: true, + }, + { + label: 'offset near end of table', + tableSize: 25, + limit: DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset: 20, + hasMore: false, + }, + { + label: 'offset exceeding table size', + tableSize: 2, + limit: DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset: 100, + hasMore: false, + }, + ])('should handle $label', async ({ tableSize, limit, offset, hasMore }) => { + const table = { ...MOCK_TABLE, - size: 100, + size: tableSize, } as unknown as DhcType.Table; - vi.mocked(mockSession.getObject).mockResolvedValue(largeTable); - - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'largeTable', - limit: 10, - }); - - expect(largeTable.setViewport).toHaveBeenCalledWith(0, 9); - expect(result.structuredContent).toMatchObject({ + vi.mocked(getTableOrError).mockResolvedValue({ success: true, - message: 'Fetched 2 rows', - details: { - hasMore: true, - limit: 10, - offset: 0, - rowCount: 2, - totalRows: 100, - }, + table, + connectionUrl: MOCK_DHC_URL, }); - }); - - it('should apply offset for pagination', async () => { - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - - const largeTable = { - ...MOCK_TABLE, - size: 100, - } as unknown as DhcType.Table; - - vi.mocked(mockSession.getObject).mockResolvedValue(largeTable); - - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'largeTable', - limit: 10, - offset: 20, + vi.mocked(getTablePage).mockResolvedValue({ + ...MOCK_PAGE_DATA, + hasMore, + totalRows: table.size, }); - expect(largeTable.setViewport).toHaveBeenCalledWith(20, 29); - expect(result.structuredContent).toMatchObject({ - success: true, - message: 'Fetched 2 rows', - details: { - hasMore: true, - limit: 10, - offset: 20, - rowCount: 2, - totalRows: 100, - }, - }); - }); - - it('should handle offset near end of table', async () => { - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - - const table = { - ...MOCK_TABLE, - size: 25, - } as unknown as DhcType.Table; - - vi.mocked(mockSession.getObject).mockResolvedValue(table); - - const tool = createGetTableDataTool({ serverManager }); + const tool = createGetTableDataTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, - tableName: 'table', - limit: 10, - offset: 20, - }); - - // Should fetch rows 20-29 (no constraint by table size) - expect(table.setViewport).toHaveBeenCalledWith(20, 29); - expect(result.structuredContent).toMatchObject({ - success: true, - message: 'Fetched 2 rows', - details: { - hasMore: false, - limit: 10, - offset: 20, - rowCount: 2, - totalRows: 25, - }, + tableName: 'testTable', + ...(limit !== DEFAULT_TABLE_PAGE_DATA_LIMIT ? { limit } : {}), + ...(offset !== DEFAULT_TABLE_PAGE_DATA_OFFSET ? { offset } : {}), }); - }); - it('should handle offset exceeding table size', async () => { - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - limit: 10, - offset: 100, + expect(getTableOrError).toHaveBeenCalledWith({ + coreJsApiCache, + connectionUrlStr: MOCK_DHC_URL.href, + variableId: undefined, + tableName: 'testTable', + serverManager, }); - // No longer returns an error - just fetches what's available + expect(getTablePage).toHaveBeenCalledWith(table, offset, limit); + expect(MOCK_TABLE.close).toHaveBeenCalled(); expect(result.structuredContent).toEqual( mcpSuccessResult('Fetched 2 rows', { - /* eslint-disable @typescript-eslint/naming-convention */ - columns: [ - { name: 'Symbol', type: 'java.lang.String' }, - { name: 'Price', type: 'double' }, - ], + ...MOCK_PAGE_DATA, connectionUrl: MOCK_DHC_URL.href, - data: [ - { Symbol: 'AAPL', Price: 150.25 }, - { Symbol: 'GOOGL', Price: 2800.5 }, - ], - hasMore: false, - limit: 10, - offset: 100, - rowCount: 2, - tableName: 'myTable', - totalRows: 2, - /* eslint-enable @typescript-eslint/naming-convention */ - }) - ); - }); - - it('should initialize session if not initialized', async () => { - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => false), - configurable: true, - }); - - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - const tool = createGetTableDataTool({ serverManager }); - await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - }); - - expect(mockConnection.getSession).toHaveBeenCalled(); - }); - - it('should handle invalid URL', async () => { - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: 'invalid-url', - tableName: 'myTable', - }); - - expect(result.structuredContent).toEqual( - mcpErrorResult('Invalid URL: Invalid URL', { - connectionUrl: 'invalid-url', - tableName: 'myTable', + hasMore, + limit, + offset, + tableName: 'testTable', + totalRows: table.size, }) ); - expect(serverManager.getServer).not.toHaveBeenCalled(); }); - it('should handle missing connection', async () => { - vi.mocked(serverManager.getServer).mockReturnValue(undefined); - - const tool = createGetTableDataTool({ serverManager }); - const result = await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', - }); - - expect(result.structuredContent).toEqual( - mcpErrorResult('No connections or server found', { + it('should propagate errors from getTableOrError', async () => { + vi.mocked(getTableOrError).mockResolvedValue({ + success: false, + errorMessage: 'Connection error', + details: { connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', - }) - ); - }); - - it('should handle missing session', async () => { - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - vi.mocked(mockConnection.getSession).mockResolvedValue(null); + }, + }); - const tool = createGetTableDataTool({ serverManager }); + const tool = createGetTableDataTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', }); - expect(result.structuredContent).toEqual( - mcpErrorResult('Unable to access session', { + expect(result.structuredContent).toMatchObject({ + success: false, + message: 'Connection error', + details: { connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', - }) - ); + }, + }); }); it('should handle errors and close table', async () => { const error = new Error('Query failed'); - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - vi.mocked(MOCK_TABLE.getViewportData).mockRejectedValue(error); + vi.mocked(getTablePage).mockRejectedValue(error); - const tool = createGetTableDataTool({ serverManager }); + const tool = createGetTableDataTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', @@ -361,11 +235,9 @@ describe('getTableData', () => { }); it('should close table on success', async () => { - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - - const tool = createGetTableDataTool({ serverManager }); + const tool = createGetTableDataTool({ coreJsApiCache, serverManager }); const result = await tool.handler({ - connectionUrl: 'http://localhost:10000', + connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', }); diff --git a/src/mcp/tools/getTableData.ts b/src/mcp/tools/getTableData.ts index 890c17e0d..c44eb23e8 100644 --- a/src/mcp/tools/getTableData.ts +++ b/src/mcp/tools/getTableData.ts @@ -1,9 +1,11 @@ import { z } from 'zod'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; import type { McpTool, McpToolHandlerArg, McpToolHandlerResult, IServerManager, + IAsyncCacheService, } from '../../types'; import { createMcpToolOutputSchema, McpToolResponse } from '../utils'; import { getTablePage, getTableOrError } from '../utils/tableUtils'; @@ -83,9 +85,14 @@ type HandlerArg = McpToolHandlerArg; type HandlerResult = McpToolHandlerResult; type GetTableDataTool = McpTool; +export const DEFAULT_TABLE_PAGE_DATA_LIMIT = 10; +export const DEFAULT_TABLE_PAGE_DATA_OFFSET = 0; + export function createGetTableDataTool({ + coreJsApiCache, serverManager, }: { + coreJsApiCache: IAsyncCacheService; serverManager: IServerManager; }): GetTableDataTool { return { @@ -93,8 +100,8 @@ export function createGetTableDataTool({ spec, handler: async ({ connectionUrl: connectionUrlStr, - limit = 10, - offset = 0, + limit = DEFAULT_TABLE_PAGE_DATA_LIMIT, + offset = DEFAULT_TABLE_PAGE_DATA_OFFSET, variableId, tableName, }: HandlerArg): Promise => { @@ -102,6 +109,7 @@ export function createGetTableDataTool({ try { const tableResult = await getTableOrError({ + coreJsApiCache, connectionUrlStr, variableId, tableName, diff --git a/src/mcp/tools/getTableStats.spec.ts b/src/mcp/tools/getTableStats.spec.ts index 114d4ddb1..fc72ff84a 100644 --- a/src/mcp/tools/getTableStats.spec.ts +++ b/src/mcp/tools/getTableStats.spec.ts @@ -2,17 +2,19 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { dh as DhcType } from '@deephaven/jsapi-types'; import { createGetTableStatsTool } from './getTableStats'; -import type { IServerManager, ServerState } from '../../types'; -import { DhcService } from '../../services'; +import type { IAsyncCacheService, IServerManager } from '../../types'; import { fakeMcpToolTimings, mcpSuccessResult, - mcpErrorResult, MOCK_DHC_URL, } from '../utils/mcpTestUtils'; +import { getTableOrError } from '../utils/tableUtils'; vi.mock('vscode'); -vi.mock('../../services/DhcService'); +vi.mock('../utils/tableUtils', () => ({ + getTableOrError: vi.fn(), + formatTableColumns: vi.fn(cols => cols), +})); const MOCK_TABLE = { size: 1000, @@ -25,23 +27,25 @@ const MOCK_TABLE = { close: vi.fn(), } as unknown as DhcType.Table; -const MOCK_SERVER_RUNNING: ServerState = { - isRunning: true, - type: 'DHC', - url: MOCK_DHC_URL, - isConnected: false, - connectionCount: 0, -}; - describe('getTableStats', () => { - const mockSession = { - getObject: vi.fn(), - } as unknown as DhcType.IdeSession; + /* eslint-disable @typescript-eslint/naming-convention */ + const mockDh = { + VariableType: { + TABLE: 'Table', + TREETABLE: 'TreeTable', + HIERARCHICALTABLE: 'HierarchicalTable', + PARTITIONEDTABLE: 'PartitionedTable', + }, + } as unknown as typeof DhcType; + /* eslint-enable @typescript-eslint/naming-convention */ - const mockConnection = Object.assign(Object.create(DhcService.prototype), { - initSession: vi.fn(), - getSession: vi.fn(), - }) as DhcService; + const coreJsApiCache = { + get: vi.fn().mockResolvedValue(mockDh), + has: vi.fn(), + invalidate: vi.fn(), + dispose: vi.fn(), + onDidInvalidate: vi.fn(), + } as unknown as IAsyncCacheService; const serverManager = { getServer: vi.fn(), @@ -53,22 +57,16 @@ describe('getTableStats', () => { vi.clearAllMocks(); fakeMcpToolTimings(); - // Default successful getServer mock - vi.mocked(serverManager.getServer).mockReturnValue(MOCK_SERVER_RUNNING); - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - // Mock isInitialized getter - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => true), - configurable: true, + // Default successful mocks + vi.mocked(getTableOrError).mockResolvedValue({ + success: true, + table: MOCK_TABLE, + connectionUrl: MOCK_DHC_URL, }); - - vi.mocked(mockConnection.getSession).mockResolvedValue(mockSession); - vi.mocked(mockSession.getObject).mockResolvedValue(MOCK_TABLE); }); it('should return correct tool spec', () => { - const tool = createGetTableStatsTool({ serverManager }); + const tool = createGetTableStatsTool({ serverManager, coreJsApiCache }); expect(tool.name).toBe('getTableStats'); expect(tool.spec.title).toBe('Get Table Schema and Statistics'); @@ -78,15 +76,18 @@ describe('getTableStats', () => { }); it('should successfully retrieve table stats', async () => { - const tool = createGetTableStatsTool({ serverManager }); + const tool = createGetTableStatsTool({ serverManager, coreJsApiCache }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', }); - expect(mockSession.getObject).toHaveBeenCalledWith({ - type: 'Table', - name: 'myTable', + expect(getTableOrError).toHaveBeenCalledWith({ + coreJsApiCache, + connectionUrlStr: MOCK_DHC_URL.href, + variableId: undefined, + tableName: 'myTable', + serverManager, }); expect(MOCK_TABLE.close).toHaveBeenCalled(); @@ -104,112 +105,34 @@ describe('getTableStats', () => { ); }); - it('should initialize session if not initialized', async () => { - Object.defineProperty(mockConnection, 'isInitialized', { - get: vi.fn(() => false), - configurable: true, - }); - - vi.mocked(serverManager.getConnections).mockReturnValue([mockConnection]); - - const tool = createGetTableStatsTool({ serverManager }); - await tool.handler({ - connectionUrl: MOCK_DHC_URL.href, - tableName: 'myTable', + it('should propagate errors from getTableOrError', async () => { + vi.mocked(getTableOrError).mockResolvedValue({ + success: false, + errorMessage: 'Connection error', + details: { + connectionUrl: MOCK_DHC_URL.href, + tableName: 'myTable', + }, }); - expect(mockConnection.getSession).toHaveBeenCalled(); - }); - - it.each([ - { - name: 'invalid URL', - connectionUrl: 'invalid-url', - server: undefined, - connections: [], - session: null, - expected: mcpErrorResult('Invalid URL: Invalid URL', { - connectionUrl: 'invalid-url', - tableName: 'mock.table', - }), - shouldCallGetServer: false, - }, - { - name: 'missing connection', - connectionUrl: MOCK_DHC_URL.href, - server: undefined, - connections: [], - session: null, - expected: mcpErrorResult('No connections or server found', { - connectionUrl: MOCK_DHC_URL.href, - tableName: 'mock.table', - }), - shouldCallGetServer: true, - }, - { - name: 'missing session', - connectionUrl: MOCK_DHC_URL.href, - server: MOCK_SERVER_RUNNING, - connections: [mockConnection], - session: null, - expected: mcpErrorResult('Unable to access session', { - connectionUrl: MOCK_DHC_URL.href, - tableName: 'mock.table', - }), - shouldCallGetServer: true, - }, - ])( - 'should handle $name', - async ({ - connectionUrl, - server, - connections, - session, - expected, - shouldCallGetServer, - }) => { - vi.mocked(serverManager.getServer).mockReturnValue(server); - vi.mocked(serverManager.getConnections).mockReturnValue(connections); - vi.mocked(mockConnection.getSession).mockResolvedValue(session); - - const tool = createGetTableStatsTool({ serverManager }); - const result = await tool.handler({ - connectionUrl, - tableName: 'mock.table', - }); - - expect(result.structuredContent).toEqual(expected); - if (!shouldCallGetServer) { - expect(serverManager.getServer).not.toHaveBeenCalled(); - } - } - ); - - it('should handle errors and close table', async () => { - const error = new Error('Table not found'); - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - vi.mocked(mockSession.getObject).mockRejectedValue(error); - - const tool = createGetTableStatsTool({ serverManager }); + const tool = createGetTableStatsTool({ serverManager, coreJsApiCache }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, - tableName: 'nonExistentTable', + tableName: 'myTable', }); expect(result.structuredContent).toMatchObject({ success: false, - message: 'Failed to get table stats: Table not found', + message: 'Connection error', details: { connectionUrl: MOCK_DHC_URL.href, - tableName: 'nonExistentTable', + tableName: 'myTable', }, }); }); it('should close table even on success', async () => { - vi.mocked(serverManager.getConnection).mockReturnValue(mockConnection); - - const tool = createGetTableStatsTool({ serverManager }); + const tool = createGetTableStatsTool({ serverManager, coreJsApiCache }); const result = await tool.handler({ connectionUrl: MOCK_DHC_URL.href, tableName: 'myTable', diff --git a/src/mcp/tools/getTableStats.ts b/src/mcp/tools/getTableStats.ts index e60675464..71f6b670c 100644 --- a/src/mcp/tools/getTableStats.ts +++ b/src/mcp/tools/getTableStats.ts @@ -1,9 +1,11 @@ import { z } from 'zod'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; import type { McpTool, McpToolHandlerArg, McpToolHandlerResult, IServerManager, + IAsyncCacheService, } from '../../types'; import { createMcpToolOutputSchema, @@ -64,7 +66,9 @@ type GetTableStatsTool = McpTool; export function createGetTableStatsTool({ serverManager, + coreJsApiCache, }: { + coreJsApiCache: IAsyncCacheService; serverManager: IServerManager; }): GetTableStatsTool { return { @@ -79,6 +83,7 @@ export function createGetTableStatsTool({ try { const tableResult = await getTableOrError({ + coreJsApiCache, connectionUrlStr, variableId, tableName, diff --git a/src/mcp/utils/mcpTestUtils.ts b/src/mcp/utils/mcpTestUtils.ts index b0f4ba8bb..6b3827fb3 100644 --- a/src/mcp/utils/mcpTestUtils.ts +++ b/src/mcp/utils/mcpTestUtils.ts @@ -89,12 +89,14 @@ export function createMockDhcService({ supportsConsoleType, getPsk, getSession, + getConnection, }: { serverUrl?: URL; runCode?: DhcType.ide.CommandResult | null; supportsConsoleType?: boolean; getPsk?: Psk | undefined; getSession?: DhcType.IdeSession | null; + getConnection?: DhcType.IdeSession | null; }): DhcService { return Object.assign(Object.create(DhcService.prototype), { serverUrl: serverUrl ?? MOCK_DHC_URL, @@ -102,5 +104,6 @@ export function createMockDhcService({ supportsConsoleType: vi.fn().mockReturnValue(supportsConsoleType ?? true), getPsk: vi.fn().mockResolvedValue(getPsk), getSession: vi.fn().mockResolvedValue(getSession ?? null), + getConnection: vi.fn().mockResolvedValue(getConnection ?? null), }); } diff --git a/src/mcp/utils/tableUtils.spec.ts b/src/mcp/utils/tableUtils.spec.ts index 5c773cf5a..5f9c58118 100644 --- a/src/mcp/utils/tableUtils.spec.ts +++ b/src/mcp/utils/tableUtils.spec.ts @@ -1,12 +1,25 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { dh as DhcType } from '@deephaven/jsapi-types'; -import type { IServerManager } from '../../types'; +import { fetchVariableDefinitionByPredicate } from '@deephaven/jsapi-utils'; +import type { IAsyncCacheService, IServerManager } from '../../types'; import { createMockDhcService, MOCK_DHC_URL } from './mcpTestUtils'; import { getFirstConnectionOrCreate } from './serverUtils'; -import { formatTableRow, formatValue, getTableOrError } from './tableUtils'; +import { + convertColumnStatsToRecords, + formatTableColumns, + formatTableRow, + formatValue, + getTableOrError, + getTablePage, + isTableType, +} from './tableUtils'; vi.mock('vscode'); +vi.mock('@deephaven/jsapi-utils', () => ({ + fetchVariableDefinitionByPredicate: vi.fn(), +})); + vi.mock('./serverUtils', () => ({ getFirstConnectionOrCreate: vi.fn(), })); @@ -19,6 +32,25 @@ describe('tableUtils', () => { describe('getTableOrError', () => { const mockTable = { type: 'Table' } as unknown as DhcType.Table; + /* eslint-disable @typescript-eslint/naming-convention */ + const mockDh = { + VariableType: { + TABLE: 'Table', + TREETABLE: 'TreeTable', + HIERARCHICALTABLE: 'HierarchicalTable', + PARTITIONEDTABLE: 'PartitionedTable', + }, + } as unknown as typeof DhcType; + /* eslint-enable @typescript-eslint/naming-convention */ + + const coreJsApiCache = { + get: vi.fn().mockResolvedValue(mockDh), + has: vi.fn(), + invalidate: vi.fn(), + dispose: vi.fn(), + onDidInvalidate: vi.fn(), + } as unknown as IAsyncCacheService; + const serverManager: IServerManager = { getServer: vi.fn(), getConnections: vi.fn(), @@ -26,9 +58,53 @@ describe('tableUtils', () => { getWorkerInfo: vi.fn(), } as unknown as IServerManager; + // Shared mock objects for reuse across tests + const mockConnection = createMockDhcService({ + serverUrl: MOCK_DHC_URL, + }); + + const mockDhConnection = { + getObject: vi.fn().mockResolvedValue(mockTable), + } as unknown as DhcType.IdeConnection; + + const mockTableVariableDef = { + type: 'Table', + id: 'mock-id', + name: 'my_table', + title: 'my_table', + } as DhcType.ide.VariableDefinition; + + const mockFigureVariableDef = { + type: 'Figure', + id: 'mock-id', + name: 'my_figure', + title: 'my_figure', + } as DhcType.ide.VariableDefinition; + + beforeEach(() => { + vi.mocked(mockConnection.getConnection).mockResolvedValue( + mockDhConnection + ); + }); + describe('error cases', () => { + it('should return error when neither variableId nor tableName provided', async () => { + const result = await getTableOrError({ + coreJsApiCache, + serverManager, + connectionUrlStr: MOCK_DHC_URL.href, + }); + + expect(result).toEqual({ + success: false, + errorMessage: 'Either variableId or tableName must be provided', + details: { connectionUrl: MOCK_DHC_URL.href }, + }); + }); + it('should return error when URL is invalid', async () => { const result = await getTableOrError({ + coreJsApiCache, serverManager, connectionUrlStr: 'not-a-valid-url', tableName: 'my_table', @@ -38,7 +114,11 @@ describe('tableUtils', () => { success: false, errorMessage: 'Invalid URL', error: expect.stringContaining('Invalid URL'), - details: { connectionUrl: 'not-a-valid-url', tableName: 'my_table' }, + details: { + connectionUrl: 'not-a-valid-url', + variableId: undefined, + tableName: 'my_table', + }, }); }); @@ -50,6 +130,7 @@ describe('tableUtils', () => { }); const result = await getTableOrError({ + coreJsApiCache, serverManager, connectionUrlStr: MOCK_DHC_URL.href, tableName: 'my_table', @@ -62,11 +143,8 @@ describe('tableUtils', () => { }); }); - it('should return error when session is not available', async () => { - const mockConnection = createMockDhcService({ - serverUrl: MOCK_DHC_URL, - }); - vi.spyOn(mockConnection, 'getSession').mockResolvedValue(null); + it('should return error when connection is not available', async () => { + vi.mocked(mockConnection.getConnection).mockResolvedValue(null); vi.mocked(getFirstConnectionOrCreate).mockResolvedValue({ success: true, @@ -75,6 +153,7 @@ describe('tableUtils', () => { }); const result = await getTableOrError({ + coreJsApiCache, serverManager, connectionUrlStr: MOCK_DHC_URL.href, tableName: 'my_table', @@ -82,21 +161,15 @@ describe('tableUtils', () => { expect(result).toEqual({ success: false, - errorMessage: 'Unable to access session', + errorMessage: 'Unable to access connection', details: { connectionUrl: MOCK_DHC_URL.href, tableName: 'my_table' }, }); }); - }); - describe('success cases', () => { - it('should return table when connection and session are available', async () => { - const mockConnection = createMockDhcService({ - serverUrl: MOCK_DHC_URL, - }); - const mockSession = { - getObject: vi.fn().mockResolvedValue(mockTable), - } as unknown as DhcType.IdeSession; - vi.spyOn(mockConnection, 'getSession').mockResolvedValue(mockSession); + it('should return error when variable is not a table type', async () => { + vi.mocked(fetchVariableDefinitionByPredicate).mockResolvedValue( + mockFigureVariableDef + ); vi.mocked(getFirstConnectionOrCreate).mockResolvedValue({ success: true, @@ -105,18 +178,75 @@ describe('tableUtils', () => { }); const result = await getTableOrError({ + coreJsApiCache, serverManager, connectionUrlStr: MOCK_DHC_URL.href, - tableName: 'my_table', + tableName: 'my_figure', }); expect(result).toEqual({ - success: true, - table: mockTable, - connectionUrl: MOCK_DHC_URL, + success: false, + errorMessage: 'Variable is not a table', + details: { + connectionUrl: MOCK_DHC_URL.href, + variableId: undefined, + tableName: 'my_figure', + }, }); }); }); + + describe('success cases', () => { + it.each(['tableName', 'variableId'])( + 'should return table when %s is provided', + async paramName => { + vi.mocked(fetchVariableDefinitionByPredicate).mockResolvedValue( + mockTableVariableDef + ); + + vi.mocked(getFirstConnectionOrCreate).mockResolvedValue({ + success: true, + connection: mockConnection, + panelUrlFormat: 'mock.panelUrlFormat', + }); + + const matchValue = 'mock.value'; + + const result = await getTableOrError({ + coreJsApiCache, + serverManager, + connectionUrlStr: MOCK_DHC_URL.href, + [paramName]: matchValue, + }); + + expect(fetchVariableDefinitionByPredicate).toHaveBeenCalledWith( + mockDhConnection, + expect.any(Function) + ); + + // Verify the predicate function works correctly + const predicate = vi.mocked(fetchVariableDefinitionByPredicate).mock + .calls[0][1]; + const predicateKey = paramName === 'tableName' ? 'name' : 'id'; + expect( + predicate({ + [predicateKey]: matchValue, + } as unknown as DhcType.ide.VariableDefinition) + ).toBe(true); + expect( + predicate({ + [predicateKey]: 'no-match', + } as unknown as DhcType.ide.VariableDefinition) + ).toBe(false); + + expect(result).toEqual({ + success: true, + table: mockTable, + connectionUrl: MOCK_DHC_URL, + }); + } + ); + }); }); describe('formatValue', () => { @@ -215,4 +345,190 @@ describe('tableUtils', () => { /* eslint-enable @typescript-eslint/naming-convention */ }); }); + + describe('convertColumnStatsToRecords', () => { + it('should convert column statistics to records', () => { + const mockStats = { + statisticsMap: new Map([ + ['MIN', 10.5], + ['MAX', 150.75], + ['AVG', 75.25], + ]), + uniqueValues: new Map([ + ['10.5', 5], + ['75.25', 10], + ['150.75', 3], + ]), + } as unknown as DhcType.ColumnStatistics; + + const result = convertColumnStatsToRecords(mockStats); + + /* eslint-disable @typescript-eslint/naming-convention */ + expect(result).toEqual({ + statistics: { + MIN: 10.5, + MAX: 150.75, + AVG: 75.25, + }, + uniqueValues: { + '10.5': 5, + '75.25': 10, + '150.75': 3, + }, + }); + /* eslint-enable @typescript-eslint/naming-convention */ + }); + + it('should handle empty statistics', () => { + const mockStats = { + statisticsMap: new Map(), + uniqueValues: new Map(), + } as unknown as DhcType.ColumnStatistics; + + const result = convertColumnStatsToRecords(mockStats); + + expect(result).toEqual({ + statistics: {}, + uniqueValues: {}, + }); + }); + }); + + describe('formatTableColumns', () => { + it('should format columns with metadata', () => { + const columns = [ + { + name: 'Symbol', + type: 'java.lang.String', + description: 'Stock symbol', + }, + { name: 'Price', type: 'double' }, + { name: 'Volume', type: 'long', description: 'Trading volume' }, + ] as DhcType.Column[]; + + const result = formatTableColumns(columns); + + expect(result).toEqual([ + { + name: 'Symbol', + type: 'java.lang.String', + description: 'Stock symbol', + }, + { name: 'Price', type: 'double' }, + { name: 'Volume', type: 'long', description: 'Trading volume' }, + ]); + }); + + it('should handle columns without descriptions', () => { + const columns = [ + { name: 'Symbol', type: 'java.lang.String', description: null }, + { name: 'Price', type: 'double' }, + ] as unknown as DhcType.Column[]; + + const result = formatTableColumns(columns); + + expect(result).toEqual([ + { name: 'Symbol', type: 'java.lang.String' }, + { name: 'Price', type: 'double' }, + ]); + }); + }); + + describe('getTablePage', () => { + it('should get paginated table data', async () => { + const mockRow1 = { + get: vi.fn((col: DhcType.Column) => { + if (col.name === 'Symbol') { + return 'AAPL'; + } + if (col.name === 'Price') { + return 150.25; + } + return null; + }), + } as unknown as DhcType.Row; + + const mockRow2 = { + get: vi.fn((col: DhcType.Column) => { + if (col.name === 'Symbol') { + return 'GOOGL'; + } + if (col.name === 'Price') { + return 2800.5; + } + return null; + }), + } as unknown as DhcType.Row; + + const mockTable = { + size: 100, + columns: [ + { name: 'Symbol', type: 'java.lang.String' }, + { name: 'Price', type: 'double' }, + ], + setViewport: vi.fn(), + getViewportData: vi.fn().mockResolvedValue({ + rows: [mockRow1, mockRow2], + }), + } as unknown as DhcType.Table; + + const result = await getTablePage(mockTable, 0, 10); + + expect(mockTable.setViewport).toHaveBeenCalledWith(0, 9); + expect(result).toEqual({ + columns: [ + { name: 'Symbol', type: 'java.lang.String' }, + { name: 'Price', type: 'double' }, + ], + data: [ + /* eslint-disable @typescript-eslint/naming-convention */ + { Symbol: 'AAPL', Price: 150.25 }, + { Symbol: 'GOOGL', Price: 2800.5 }, + /* eslint-enable @typescript-eslint/naming-convention */ + ], + hasMore: true, + rowCount: 2, + totalRows: 100, + }); + }); + + it('should set hasMore to false when at end of table', async () => { + const mockTable = { + size: 5, + columns: [{ name: 'Symbol', type: 'java.lang.String' }], + setViewport: vi.fn(), + getViewportData: vi.fn().mockResolvedValue({ rows: [] }), + } as unknown as DhcType.Table; + + const result = await getTablePage(mockTable, 0, 10); + + expect(result.hasMore).toBe(false); + }); + }); + + describe('isTableType', () => { + /* eslint-disable @typescript-eslint/naming-convention */ + const mockDh = { + VariableType: { + TABLE: 'Table', + TREETABLE: 'TreeTable', + HIERARCHICALTABLE: 'HierarchicalTable', + PARTITIONEDTABLE: 'PartitionedTable', + FIGURE: 'Figure', + }, + } as unknown as typeof DhcType; + /* eslint-enable @typescript-eslint/naming-convention */ + + it.each([ + { type: 'Table', expected: true }, + { type: 'TreeTable', expected: true }, + { type: 'HierarchicalTable', expected: true }, + { type: 'PartitionedTable', expected: true }, + { type: 'Figure', expected: false }, + { type: 'OtherType', expected: false }, + ])('should return $expected for type $type', ({ type, expected }) => { + const result = isTableType(mockDh, type); + expect(result).toBe(expected); + }); + }); }); diff --git a/src/mcp/utils/tableUtils.ts b/src/mcp/utils/tableUtils.ts index c56c6d2b8..a0a5cf7c4 100644 --- a/src/mcp/utils/tableUtils.ts +++ b/src/mcp/utils/tableUtils.ts @@ -1,5 +1,6 @@ import type { dh as DhType } from '@deephaven/jsapi-types'; -import type { IServerManager } from '../../types'; +import { fetchVariableDefinitionByPredicate } from '@deephaven/jsapi-utils'; +import type { IAsyncCacheService, IServerManager } from '../../types'; import { parseUrl } from '../../util'; import { getFirstConnectionOrCreate } from './serverUtils'; @@ -112,19 +113,27 @@ export function formatValue(value: unknown): unknown { * 6. Fetching the table by name * * @param params Configuration for getting the table - * @param params.serverManager The server manager to query * @param params.connectionUrlStr The connection URL string + * @param params.coreJsApiCache Cache for the DH JS API + * @param params.serverManager The server manager to query * @param params.variableId Optional variable ID to fetch (takes precedence over tableName if provided) * @param params.tableName Optional name of the table to fetch (if variableId is not provided) * @returns Success with table, or error with message, hint, and details */ export async function getTableOrError(params: { - serverManager: IServerManager; connectionUrlStr: string; + coreJsApiCache: IAsyncCacheService; + serverManager: IServerManager; variableId?: string; tableName?: string; }): Promise { - const { serverManager, connectionUrlStr, variableId, tableName } = params; + const { + coreJsApiCache, + connectionUrlStr, + serverManager, + variableId, + tableName, + } = params; if (variableId == null && tableName == null) { return { @@ -157,27 +166,35 @@ export async function getTableOrError(params: { } const { connection } = firstConnectionResult; - const session = await connection.getSession(); + const cn = await connection.getConnection(); - if (session == null) { + if (cn == null) { return { success: false, - errorMessage: 'Unable to access session', + errorMessage: 'Unable to access connection', details: { connectionUrl: connectionUrlStr, tableName }, }; } - const table = await session.getObject( - variableId == null - ? { - type: 'Table', - name: tableName, - } - : { - type: 'Table', - id: variableId, - } - ); + const variableDef: DhType.ide.VariableDefinition = + await fetchVariableDefinitionByPredicate( + cn, + variableId == null + ? (def): boolean => def.name === tableName + : (def): boolean => def.id === variableId + ); + + const dh = await coreJsApiCache.get(connection.serverUrl); + + if (!isTableType(dh, variableDef.type)) { + return { + success: false, + errorMessage: 'Variable is not a table', + details: { connectionUrl: connectionUrlStr, variableId, tableName }, + }; + } + + const table = await cn.getObject(variableDef); return { success: true, @@ -222,3 +239,18 @@ export async function getTablePage( totalRows, }; } + +/** + * Check if a given variable type is a table type. + * @param dh The DH JS API instance + * @param type The variable type to check + * @returns True if the type is a table type, false otherwise + */ +export function isTableType(dh: typeof DhType, type: string): boolean { + return ( + type === dh.VariableType.TABLE || + type === dh.VariableType.TREETABLE || + type === dh.VariableType.HIERARCHICALTABLE || + type === dh.VariableType.PARTITIONEDTABLE + ); +} diff --git a/src/services/DhcService.ts b/src/services/DhcService.ts index e6b55ba8b..edfa04676 100644 --- a/src/services/DhcService.ts +++ b/src/services/DhcService.ts @@ -405,6 +405,14 @@ export class DhcService extends DisposableBase implements IDhcService { await saveRequirementsTxt(dependencies); } + async getConnection(): Promise { + if (this.cn == null) { + await this.initSession(); + } + + return this.cn; + } + async getSession(): Promise { if (this.session == null) { await this.initSession(); diff --git a/src/services/ServerManager.ts b/src/services/ServerManager.ts index c4ade86c1..334b5029c 100644 --- a/src/services/ServerManager.ts +++ b/src/services/ServerManager.ts @@ -52,6 +52,7 @@ export class ServerManager implements IServerManager { ) { this._configService = configService; this._connectionMap = new URLMap(); + this._pendingConnectionMap = new URLMap>(); this._coreClientCache = coreClientCache; this._dhcServiceFactory = dhcServiceFactory; this._dheClientCache = dheClientCache; @@ -70,6 +71,9 @@ export class ServerManager implements IServerManager { private readonly _configService: IConfigService; private readonly _connectionMap: URLMap; + private readonly _pendingConnectionMap: URLMap< + Promise + >; private readonly _coreClientCache: URLMap; private readonly _dhcServiceFactory: IDhcServiceFactory; private readonly _dheClientCache: URLMap; @@ -115,6 +119,7 @@ export class ServerManager implements IServerManager { await Promise.all([ this._connectionMap.dispose(), + this._pendingConnectionMap.dispose(), this._serverMap.dispose(), this._uriConnectionsMap.dispose(), this._workerURLToServerURLMap.dispose(), @@ -200,18 +205,52 @@ export class ServerManager implements IServerManager { const serverState = this._serverMap.get(serverUrl); if (serverState == null) { - return null; + throw new Error(`Server with URL '${serverUrl}' not found.`); } - // DHE supports multiple connections, but DHC does not. - if ( - !this._dheServiceCache.has(serverUrl) && - serverState.connectionCount > 0 - ) { - logger.info('Already connected to server:', serverUrl); - return null; + // We only support 1 connection for DHC servers in the extension + if (serverState.type === 'DHC' && serverState.connectionCount > 0) { + logger.info('Already connected to server:', serverUrl.href); + return this._connectionMap.getOrThrow(serverUrl); + } + + if (this._pendingConnectionMap.has(serverUrl)) { + logger.debug('Connection already in progress:', serverUrl.href); + return this._pendingConnectionMap.getOrThrow(serverUrl); + } + + const connectionPromise = this._doConnectToServer( + serverState, + workerConsoleType, + operateAsAnotherUser + ); + + // We only support 1 connection for DHC servers in the extension, but the + // count doesn't get updated until the connection is established, so we need + // to mark pending connections to prevent multiple simultaneous connection + // attempts to the same DHC server. + if (serverState.type === 'DHC') { + this._pendingConnectionMap.set( + serverUrl, + connectionPromise.then(result => { + this._pendingConnectionMap.delete(serverUrl); + return result; + }) + ); } + return connectionPromise; + }; + + private _doConnectToServer = async ( + serverState: ServerState, + workerConsoleType?: ConsoleType, + operateAsAnotherUser: boolean = false + ): Promise => { + let serverUrl = serverState.url; + + logger.debug('Connecting to server:', serverUrl.href); + let tagId: UniqueID | undefined; let placeholderUrl: URL | undefined; diff --git a/src/types/serviceTypes.d.ts b/src/types/serviceTypes.d.ts index 194664fd1..5f056c02f 100644 --- a/src/types/serviceTypes.d.ts +++ b/src/types/serviceTypes.d.ts @@ -63,6 +63,7 @@ export interface IDhcService extends IDisposable, ConnectionState { initSession(): Promise; getClient(): Promise; + getConnection(): Promise; getConsoleTypes: () => Promise>; getPsk(): Promise; getSession(): Promise;