diff --git a/README.md b/README.md index 7ae91e6..00d4e15 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ your IB account to retrieve market data, check positions, and place trades. ## Features - **Interactive Brokers API Integration**: Full trading capabilities including account management, position tracking, real-time market data, and order management (market, limit, and stop orders) +- **Flex Query Support**: Execute Flex Queries to retrieve account statements, trade confirmations, and historical data. Queries are automatically remembered for easy reuse - **Flexible Authentication**: Choose between browser-based OAuth authentication or headless mode with credentials for automated environments - **Simple Setup**: Run directly with `npx` - no Docker or additional installations required. Includes pre-configured IB Gateway and Java runtime for all platforms @@ -125,6 +126,50 @@ To enable paper trading, add `"IB_PAPER_TRADING": "true"` to your environment va control. Consider using environment variable files or secure credential management systems. +## Flex Query Configuration (Optional) + +To use Flex Queries for retrieving account statements and historical data, you need to configure your Flex Web Service Token: + +```json +{ + "mcpServers": { + "interactive-brokers": { + "command": "npx", + "args": ["-y", "interactive-brokers-mcp"], + "env": { + "IB_FLEX_TOKEN": "your_flex_token_here" + } + } + } +} +``` + +### How to Get Your Flex Token: + +1. Log in to [Interactive Brokers Account Management](https://www.interactivebrokers.com/portal) +2. Go to **Settings** → **Account Settings** +3. Navigate to **Reporting** → **Flex Web Service** +4. Generate or retrieve your Flex Web Service Token + +For detailed instructions on enabling Flex Web Service, see the [IB Flex Web Service Guide](https://www.ibkrguides.com/orgportal/performanceandstatements/flex-web-service.htm). + +### Creating Flex Queries: + +1. Go to **Reports** → **Flex Queries** in Account Management +2. Create or customize your query template +3. Click the info icon next to your query to find its Query ID + +For a complete guide on creating and customizing Flex Queries, see the [IB Flex Queries Guide](https://www.ibkrguides.com/orgportal/performanceandstatements/flex.htm). + +**Note**: When you execute a Flex Query for the first time, the MCP server automatically saves it with its name from the API. Future executions can reference the query by either its ID or its saved name. + +### Flex Query Features: + +- **Automatic Memory**: When you execute a Flex Query, it's automatically saved for future use +- **Easy Reuse**: Previously used queries are remembered - no need to copy query IDs repeatedly +- **Friendly Names**: Optionally provide a friendly name when first executing a query +- **Forget Queries**: Remove queries you no longer need with the `forget_flex_query` tool + ## Configuration Variables | Feature | Environment Variable | Command Line Argument | @@ -134,9 +179,12 @@ management systems. | Headless Mode | `IB_HEADLESS_MODE` | `--ib-headless-mode` | | Paper Trading | `IB_PAPER_TRADING` | `--ib-paper-trading` | | Auth Timeout | `IB_AUTH_TIMEOUT` | `--ib-auth-timeout` | +| Flex Token | `IB_FLEX_TOKEN` | N/A | ## Available MCP Tools +### Trading & Account Management + | Tool | Description | | ------------------ | ----------------------------------------- | | `get_account_info` | Retrieve account information and balances | @@ -145,10 +193,14 @@ management systems. | `place_order` | Place market, limit, or stop orders | | `get_order_status` | Check order execution status | | `get_live_orders` | Get all live/open orders for monitoring | -| `get_alerts` | Get all trading alerts for an account | -| `create_alert` | Create a new price or condition alert | -| `activate_alert` | Activate a previously created alert | -| `delete_alert` | Delete an existing alert | + +### Flex Queries (Requires IB_FLEX_TOKEN) + +| Tool | Description | +| ------------------- | -------------------------------------------------------------------- | +| `get_flex_query` | Execute a Flex Query and retrieve statements (auto-saves for reuse) | +| `list_flex_queries` | List all previously used Flex Queries | +| `forget_flex_query` | Remove a saved Flex Query from memory | ## Troubleshooting diff --git a/mcp-inspector.json b/mcp-inspector.json index e52b270..a1d0e06 100644 --- a/mcp-inspector.json +++ b/mcp-inspector.json @@ -1,6 +1,6 @@ { "mcpServers": { - "interactive-brokers": { + "interactive-brokers-mcp": { "type": "sse", "url": "http://localhost:8123/mcp" } diff --git a/package-lock.json b/package-lock.json index 0c3b651..6fe0fef 100644 --- a/package-lock.json +++ b/package-lock.json @@ -16,6 +16,7 @@ "express": "^4.18.0", "open": "^10.2.0", "playwright-core": "^1.54.2", + "xml2js": "^0.6.2", "zod": "^3.22.0" }, "bin": { @@ -29,6 +30,7 @@ "@types/cors": "^2.8.0", "@types/express": "^4.17.0", "@types/node": "^20.19.22", + "@types/xml2js": "^0.4.14", "@vitest/coverage-v8": "^3.2.4", "@vitest/ui": "^3.2.4", "conventional-changelog-conventionalcommits": "^9.1.0", @@ -1999,6 +2001,16 @@ "@types/send": "*" } }, + "node_modules/@types/xml2js": { + "version": "0.4.14", + "resolved": "https://registry.npmjs.org/@types/xml2js/-/xml2js-0.4.14.tgz", + "integrity": "sha512-4YnrRemBShWRO2QjvUin8ESA41rH+9nQGLUGZV/1IDhi3SL9OhdpNC/MrulTWuptXKwhx/aDxE7toV0f/ypIXQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@vitest/coverage-v8": { "version": "3.2.4", "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.4.tgz", @@ -8581,6 +8593,12 @@ "integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==", "license": "MIT" }, + "node_modules/sax": { + "version": "1.4.2", + "resolved": "https://registry.npmjs.org/sax/-/sax-1.4.2.tgz", + "integrity": "sha512-FySGAa0RGcFiN6zfrO9JvK1r7TB59xuzCcTHOBXBNoKgDejlOQCR2KL/FGk3/iDlsqyYg1ELZpOmlg09B01Czw==", + "license": "BlueOak-1.0.0" + }, "node_modules/semantic-release": { "version": "22.0.12", "resolved": "https://registry.npmjs.org/semantic-release/-/semantic-release-22.0.12.tgz", @@ -10078,6 +10096,28 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/xml2js": { + "version": "0.6.2", + "resolved": "https://registry.npmjs.org/xml2js/-/xml2js-0.6.2.tgz", + "integrity": "sha512-T4rieHaC1EXcES0Kxxj4JWgaUQHDk+qwHcYOCFHfiwKz7tOVPLq7Hjq9dM1WCMhylqMEfP7hMcOIChvotiZegA==", + "license": "MIT", + "dependencies": { + "sax": ">=0.6.0", + "xmlbuilder": "~11.0.0" + }, + "engines": { + "node": ">=4.0.0" + } + }, + "node_modules/xmlbuilder": { + "version": "11.0.1", + "resolved": "https://registry.npmjs.org/xmlbuilder/-/xmlbuilder-11.0.1.tgz", + "integrity": "sha512-fDlsI/kFEx7gLvbecc0/ohLG50fugQp8ryHzMTuW9vSa1GJ0XYWKnhsUx7oie3G98+r56aTQIUB4kht42R3JvA==", + "license": "MIT", + "engines": { + "node": ">=4.0" + } + }, "node_modules/xtend": { "version": "4.0.2", "resolved": "https://registry.npmjs.org/xtend/-/xtend-4.0.2.tgz", diff --git a/package.json b/package.json index 978b56f..9b315b1 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,8 @@ "test": "vitest run", "test:watch": "vitest", "test:ui": "vitest --ui", - "test:coverage": "vitest run --coverage" + "test:coverage": "vitest run --coverage", + "inspector": "npx @modelcontextprotocol/inspector --config mcp-inspector.json --server interactive-brokers-mcp" }, "keywords": [ "mcp", @@ -46,6 +47,7 @@ "express": "^4.18.0", "open": "^10.2.0", "playwright-core": "^1.54.2", + "xml2js": "^0.6.2", "zod": "^3.22.0" }, "devDependencies": { @@ -56,6 +58,7 @@ "@types/cors": "^2.8.0", "@types/express": "^4.17.0", "@types/node": "^20.19.22", + "@types/xml2js": "^0.4.14", "@vitest/coverage-v8": "^3.2.4", "@vitest/ui": "^3.2.4", "conventional-changelog-conventionalcommits": "^9.1.0", diff --git a/src/config.ts b/src/config.ts index e116b96..988407c 100644 --- a/src/config.ts +++ b/src/config.ts @@ -17,5 +17,8 @@ export const config = { // Paper trading configuration IB_PAPER_TRADING: process.env.IB_PAPER_TRADING === "true", + + // Flex Query configuration + IB_FLEX_TOKEN: process.env.IB_FLEX_TOKEN || "", }; \ No newline at end of file diff --git a/src/flex-query-client.ts b/src/flex-query-client.ts new file mode 100644 index 0000000..044b00d --- /dev/null +++ b/src/flex-query-client.ts @@ -0,0 +1,228 @@ +import axios, { AxiosInstance } from "axios"; +import { Logger } from "./logger.js"; +import { parseStringPromise } from "xml2js"; + +export interface FlexQueryClientConfig { + token: string; +} + +export interface FlexQueryResponse { + referenceCode?: string; + url?: string; + error?: string; + errorCode?: string; +} + +export interface FlexStatementResponse { + data?: string; // XML data + error?: string; + errorCode?: string; +} + +/** + * Client for Interactive Brokers Flex Query Web Service + * API Documentation: https://www.interactivebrokers.com/en/software/am/am/reports/flex_web_service_version_3.htm + */ +export class FlexQueryClient { + private client: AxiosInstance; + private token: string; + private baseUrl = "https://gdcdyn.interactivebrokers.com/Universal/servlet"; + + constructor(config: FlexQueryClientConfig) { + this.token = config.token; + this.client = axios.create({ + timeout: 60000, // Flex queries can take a while + }); + } + + /** + * Send a request to execute a flex query + * @param queryId The flex query ID to execute + * @returns Response containing reference code or error + */ + async sendRequest(queryId: string): Promise { + try { + Logger.log(`[FLEX-QUERY] Sending request for query ID: ${queryId}`); + + const url = `${this.baseUrl}/FlexStatementService.SendRequest`; + const params = { + t: this.token, + q: queryId, + v: "3", // API version + }; + + const response = await this.client.get(url, { params }); + + Logger.log(`[FLEX-QUERY] SendRequest response:`, response.data); + + // Parse XML response + const parsed = await parseStringPromise(response.data, { explicitArray: false }); + + if (parsed.FlexStatementResponse) { + const flexResponse = parsed.FlexStatementResponse; + + if (flexResponse.Status === "Success") { + return { + referenceCode: flexResponse.ReferenceCode, + url: flexResponse.Url, + }; + } else if (flexResponse.Status === "Fail") { + return { + error: flexResponse.ErrorMessage || flexResponse.ErrorCode || "Unknown error", + errorCode: flexResponse.ErrorCode, + }; + } + } + + throw new Error("Unexpected response format from Flex Query service"); + } catch (error) { + Logger.error("[FLEX-QUERY] Failed to send request:", error); + + if (axios.isAxiosError(error)) { + throw new Error(`Failed to send flex query request: ${error.message}`); + } + + throw error; + } + } + + /** + * Get the statement data using a reference code + * @param referenceCode The reference code from sendRequest + * @returns The flex statement data (XML format) or error + */ + async getStatement(referenceCode: string): Promise { + try { + Logger.log(`[FLEX-QUERY] Getting statement for reference code: ${referenceCode}`); + + const url = `${this.baseUrl}/FlexStatementService.GetStatement`; + const params = { + t: this.token, + q: referenceCode, + v: "3", // API version + }; + + const response = await this.client.get(url, { params }); + + // Parse XML response + const parsed = await parseStringPromise(response.data, { explicitArray: false }); + + if (parsed.FlexStatementResponse) { + const flexResponse = parsed.FlexStatementResponse; + + if (flexResponse.Status === "Success") { + // The actual statement data is in the response + return { + data: response.data, + }; + } else if (flexResponse.Status === "Fail") { + return { + error: flexResponse.ErrorMessage || flexResponse.ErrorCode || "Unknown error", + errorCode: flexResponse.ErrorCode, + }; + } + } else if (parsed.FlexQueryResponse) { + // This is the actual statement data + return { + data: response.data, + }; + } + + throw new Error("Unexpected response format from Flex Query service"); + } catch (error) { + Logger.error("[FLEX-QUERY] Failed to get statement:", error); + + if (axios.isAxiosError(error)) { + throw new Error(`Failed to get flex statement: ${error.message}`); + } + + throw error; + } + } + + /** + * Execute a flex query and wait for the results + * @param queryId The flex query ID to execute + * @param maxRetries Maximum number of retries for getting the statement (default: 10) + * @param retryDelayMs Delay between retries in milliseconds (default: 2000) + * @returns The flex statement data or error + */ + async executeQuery( + queryId: string, + maxRetries: number = 10, + retryDelayMs: number = 2000 + ): Promise { + // Step 1: Send the request + const sendResponse = await this.sendRequest(queryId); + + if (sendResponse.error) { + return { + error: sendResponse.error, + errorCode: sendResponse.errorCode, + }; + } + + if (!sendResponse.referenceCode) { + return { + error: "No reference code received from flex query service", + }; + } + + Logger.log(`[FLEX-QUERY] Query submitted, reference code: ${sendResponse.referenceCode}`); + Logger.log(`[FLEX-QUERY] Waiting for statement to be ready (max ${maxRetries} retries)...`); + + // Step 2: Poll for the statement + for (let i = 0; i < maxRetries; i++) { + // Wait before checking + await new Promise((resolve) => setTimeout(resolve, retryDelayMs)); + + Logger.log(`[FLEX-QUERY] Attempt ${i + 1}/${maxRetries} to retrieve statement...`); + + const statementResponse = await this.getStatement(sendResponse.referenceCode); + + if (statementResponse.error) { + // Check if it's a "not ready yet" error + if ( + statementResponse.errorCode === "1019" || // Statement generation in progress + statementResponse.error.includes("in progress") || + statementResponse.error.includes("not ready") + ) { + Logger.log(`[FLEX-QUERY] Statement not ready yet, retrying...`); + continue; + } + + // It's a real error + return statementResponse; + } + + // Success! + Logger.log(`[FLEX-QUERY] Statement retrieved successfully`); + return statementResponse; + } + + return { + error: `Statement not ready after ${maxRetries} retries. Please try again later.`, + }; + } + + /** + * Parse flex statement XML data into a more usable JSON format + * @param xmlData The XML data from getStatement + * @returns Parsed JSON object + */ + async parseStatement(xmlData: string): Promise { + try { + const parsed = await parseStringPromise(xmlData, { + explicitArray: false, + mergeAttrs: true, + }); + return parsed; + } catch (error) { + Logger.error("[FLEX-QUERY] Failed to parse statement:", error); + throw new Error("Failed to parse flex statement XML"); + } + } +} + + + diff --git a/src/flex-query-storage.ts b/src/flex-query-storage.ts new file mode 100644 index 0000000..5d00a92 --- /dev/null +++ b/src/flex-query-storage.ts @@ -0,0 +1,235 @@ +import fs from "fs/promises"; +import path from "path"; +import { Logger } from "./logger.js"; +import os from "os"; + +export interface SavedFlexQuery { + id: string; + name: string; + queryId: string; + description?: string; + createdAt: string; + lastUsed?: string; +} + +export interface FlexQueriesStore { + queries: SavedFlexQuery[]; +} + +/** + * Storage manager for saved flex queries + * Uses a JSON file to persist query information + */ +export class FlexQueryStorage { + private storageFile: string; + private store: FlexQueriesStore = { queries: [] }; + + constructor(storageFile?: string) { + // Use user's home directory for storage by default + const defaultFile = path.join( + os.homedir(), + ".interactive-brokers-mcp", + "flex-queries.json" + ); + this.storageFile = storageFile || defaultFile; + } + + /** + * Initialize storage - create directory and file if needed + */ + async initialize(): Promise { + try { + const dir = path.dirname(this.storageFile); + + // Create directory if it doesn't exist + try { + await fs.access(dir); + } catch { + Logger.log(`[FLEX-QUERY-STORAGE] Creating directory: ${dir}`); + await fs.mkdir(dir, { recursive: true }); + } + + // Load existing data or create new file + try { + await this.load(); + } catch { + Logger.log(`[FLEX-QUERY-STORAGE] Creating new storage file: ${this.storageFile}`); + await this.save(); + } + } catch (error) { + Logger.error("[FLEX-QUERY-STORAGE] Failed to initialize storage:", error); + throw error; + } + } + + /** + * Load queries from storage file + */ + private async load(): Promise { + try { + const data = await fs.readFile(this.storageFile, "utf-8"); + this.store = JSON.parse(data); + Logger.log(`[FLEX-QUERY-STORAGE] Loaded ${this.store.queries.length} queries from storage`); + } catch (error) { + if ((error as any).code === "ENOENT") { + // File doesn't exist yet, use empty store + this.store = { queries: [] }; + return; + } + throw error; + } + } + + /** + * Save queries to storage file + */ + private async save(): Promise { + try { + const data = JSON.stringify(this.store, null, 2); + await fs.writeFile(this.storageFile, data, "utf-8"); + Logger.log(`[FLEX-QUERY-STORAGE] Saved ${this.store.queries.length} queries to storage`); + } catch (error) { + Logger.error("[FLEX-QUERY-STORAGE] Failed to save queries:", error); + throw error; + } + } + + /** + * Get all saved queries + */ + async listQueries(): Promise { + await this.load(); // Reload to get latest data + return [...this.store.queries]; + } + + /** + * Get a saved query by internal ID + */ + async getQuery(id: string): Promise { + await this.load(); + const query = this.store.queries.find((q) => q.id === id); + return query || null; + } + + /** + * Get a saved query by IB's query ID + */ + async getQueryByQueryId(queryId: string): Promise { + await this.load(); + const query = this.store.queries.find((q) => q.queryId === queryId); + return query || null; + } + + /** + * Get a saved query by name + */ + async getQueryByName(name: string): Promise { + await this.load(); + const query = this.store.queries.find((q) => q.name.toLowerCase() === name.toLowerCase()); + return query || null; + } + + /** + * Save a new query + */ + async saveQuery(query: Omit): Promise { + await this.load(); + + // Check if a query with this queryId (IB's ID) already exists + const existing = this.store.queries.find((q) => q.queryId === query.queryId); + if (existing) { + throw new Error(`A query with the IB Query ID "${query.queryId}" already exists. Use updateQuery to modify it.`); + } + + const newQuery: SavedFlexQuery = { + ...query, + id: this.generateId(), + createdAt: new Date().toISOString(), + }; + + this.store.queries.push(newQuery); + await this.save(); + + Logger.log(`[FLEX-QUERY-STORAGE] Saved new query: ${newQuery.name} (Query ID: ${newQuery.queryId})`); + return newQuery; + } + + /** + * Update an existing query + */ + async updateQuery(id: string, updates: Partial>): Promise { + await this.load(); + + const index = this.store.queries.findIndex((q) => q.id === id); + if (index === -1) { + throw new Error(`Query with ID "${id}" not found`); + } + + // If updating name, check for conflicts + if (updates.name) { + const nameConflict = this.store.queries.find( + (q) => q.id !== id && q.name.toLowerCase() === updates.name!.toLowerCase() + ); + if (nameConflict) { + throw new Error(`A query with the name "${updates.name}" already exists`); + } + } + + this.store.queries[index] = { + ...this.store.queries[index], + ...updates, + }; + + await this.save(); + + Logger.log(`[FLEX-QUERY-STORAGE] Updated query: ${this.store.queries[index].name} (ID: ${id})`); + return this.store.queries[index]; + } + + /** + * Update the last used timestamp for a query + */ + async markQueryUsed(id: string): Promise { + await this.load(); + + const query = this.store.queries.find((q) => q.id === id); + if (query) { + query.lastUsed = new Date().toISOString(); + await this.save(); + } + } + + /** + * Delete a saved query + */ + async deleteQuery(id: string): Promise { + await this.load(); + + const index = this.store.queries.findIndex((q) => q.id === id); + if (index === -1) { + return false; + } + + const deleted = this.store.queries.splice(index, 1)[0]; + await this.save(); + + Logger.log(`[FLEX-QUERY-STORAGE] Deleted query: ${deleted.name} (ID: ${id})`); + return true; + } + + /** + * Generate a unique ID for a query + */ + private generateId(): string { + return `query_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; + } + + /** + * Get the storage file path + */ + getStorageFilePath(): string { + return this.storageFile; + } +} + + diff --git a/src/server.ts b/src/server.ts index 5c308d9..6fd17ad 100644 --- a/src/server.ts +++ b/src/server.ts @@ -86,3 +86,8 @@ export function createIBMCPServer({ config: userConfig }: { config: z.infer; export type GetAccountInfoInput = z.infer; @@ -145,3 +167,6 @@ export type GetAlertsInput = z.infer; export type CreateAlertInput = z.infer; export type ActivateAlertInput = z.infer; export type DeleteAlertInput = z.infer; +export type GetFlexQueryInput = z.infer; +export type ListFlexQueriesInput = z.infer; +export type ForgetFlexQueryInput = z.infer; diff --git a/src/tool-handlers.ts b/src/tool-handlers.ts index 20df772..edb1c0f 100644 --- a/src/tool-handlers.ts +++ b/src/tool-handlers.ts @@ -3,6 +3,8 @@ import { IBGatewayManager } from "./gateway-manager.js"; import { HeadlessAuthenticator, HeadlessAuthConfig } from "./headless-auth.js"; import open from "open"; import { Logger } from "./logger.js"; +import { FlexQueryClient } from "./flex-query-client.js"; +import { FlexQueryStorage } from "./flex-query-storage.js"; import { AuthenticateInput, GetAccountInfoInput, @@ -16,12 +18,17 @@ import { CreateAlertInput, ActivateAlertInput, DeleteAlertInput, + GetFlexQueryInput, + ListFlexQueriesInput, + ForgetFlexQueryInput, } from "./tool-definitions.js"; export interface ToolHandlerContext { ibClient: IBClient; gatewayManager?: IBGatewayManager; config: any; + flexQueryClient?: FlexQueryClient; + flexQueryStorage?: FlexQueryStorage; } export type ToolHandlerResult = { @@ -36,6 +43,22 @@ export class ToolHandlers { constructor(context: ToolHandlerContext) { this.context = context; + + // Initialize flex query client and storage if token is provided + // Only initialize if not already set (useful for testing) + if (context.config.IB_FLEX_TOKEN && !context.flexQueryClient) { + this.context.flexQueryClient = new FlexQueryClient({ + token: context.config.IB_FLEX_TOKEN, + }); + } + + if (context.config.IB_FLEX_TOKEN && !context.flexQueryStorage) { + this.context.flexQueryStorage = new FlexQueryStorage(); + // Initialize storage asynchronously + this.context.flexQueryStorage.initialize().catch((error) => { + Logger.error("[FLEX-QUERY] Failed to initialize storage:", error); + }); + } } // Ensure Gateway is ready before operations @@ -635,4 +658,242 @@ export class ToolHandlers { }; } } + + // ── Flex Query Methods ────────────────────────────────────────────────────── + + async getFlexQuery(input: GetFlexQueryInput): Promise { + try { + if (!this.context.flexQueryClient) { + return { + content: [ + { + type: "text", + text: JSON.stringify({ + error: "Flex Query feature not configured", + message: "Please set the IB_FLEX_TOKEN environment variable to use Flex Queries", + instructions: [ + "1. Get your Flex Web Service Token from Interactive Brokers", + "2. Set the IB_FLEX_TOKEN environment variable", + "3. Restart the MCP server" + ] + }, null, 2), + }, + ], + }; + } + + if (!this.context.flexQueryStorage) { + throw new Error("Flex Query storage not initialized"); + } + + Logger.log(`[FLEX-QUERY] Executing flex query: ${input.queryId}`); + + // Check if this query was used before (by IB's query ID) + const existingQuery = await this.context.flexQueryStorage.getQueryByQueryId(input.queryId); + + // Execute the query + const result = await this.context.flexQueryClient.executeQuery(input.queryId); + + if (result.error) { + return { + content: [ + { + type: "text", + text: JSON.stringify({ + error: result.error, + errorCode: result.errorCode, + queryId: input.queryId, + }, null, 2), + }, + ], + }; + } + + // Parse XML to extract query name from the response + let parsedData; + let queryNameFromApi: string | undefined; + + if (result.data) { + try { + parsedData = await this.context.flexQueryClient.parseStatement(result.data); + + // Extract query name from the parsed XML + // The queryName is directly under FlexQueryResponse + if (parsedData?.FlexQueryResponse) { + queryNameFromApi = parsedData.FlexQueryResponse.queryName; + } + + Logger.log(`[FLEX-QUERY] Extracted query name from API: ${queryNameFromApi}`); + } catch (parseError) { + Logger.warn("[FLEX-QUERY] Failed to parse XML for query name extraction:", parseError); + } + } + + // Auto-save the query if it's new or update last used + if (existingQuery) { + await this.context.flexQueryStorage.markQueryUsed(existingQuery.id); + Logger.log(`[FLEX-QUERY] Updated last used timestamp for query: ${input.queryId}`); + } else { + // Save new query with the name from API, input, or fallback to queryId + const queryName = queryNameFromApi || input.queryName || input.queryId; + await this.context.flexQueryStorage.saveQuery({ + name: queryName, + queryId: input.queryId, + description: `Auto-saved on ${new Date().toLocaleDateString()}`, + }); + Logger.log(`[FLEX-QUERY] Auto-saved new query: ${queryName}`); + } + + // Return parsed data if requested (and we haven't parsed it yet) + if (input.parseXml && !parsedData && result.data) { + try { + parsedData = await this.context.flexQueryClient.parseStatement(result.data); + } catch (parseError) { + Logger.warn("[FLEX-QUERY] Failed to parse XML, returning raw data:", parseError); + } + } + + return { + content: [ + { + type: "text", + text: JSON.stringify({ + success: true, + queryId: input.queryId, + queryName: queryNameFromApi, + autoSaved: !existingQuery, + data: parsedData || result.data, + note: existingQuery + ? "Query was previously saved and has been marked as used" + : "Query has been automatically saved for future reference" + }, null, 2), + }, + ], + }; + } catch (error) { + return { + content: [ + { + type: "text", + text: this.formatError(error), + }, + ], + }; + } + } + + async listFlexQueries(input: ListFlexQueriesInput): Promise { + try { + if (!this.context.flexQueryStorage) { + return { + content: [ + { + type: "text", + text: JSON.stringify({ + error: "Flex Query feature not configured", + message: "Please set the IB_FLEX_TOKEN environment variable to use Flex Queries" + }, null, 2), + }, + ], + }; + } + + const queries = await this.context.flexQueryStorage.listQueries(); + + return { + content: [ + { + type: "text", + text: JSON.stringify({ + count: queries.length, + queries: queries.map(q => ({ + name: q.name, + queryId: q.queryId, + description: q.description, + createdAt: q.createdAt, + lastUsed: q.lastUsed, + })), + storageLocation: this.context.flexQueryStorage.getStorageFilePath(), + }, null, 2), + }, + ], + }; + } catch (error) { + return { + content: [ + { + type: "text", + text: this.formatError(error), + }, + ], + }; + } + } + + async forgetFlexQuery(input: ForgetFlexQueryInput): Promise { + try { + if (!this.context.flexQueryStorage) { + return { + content: [ + { + type: "text", + text: JSON.stringify({ + error: "Flex Query feature not configured", + message: "Please set the IB_FLEX_TOKEN environment variable to use Flex Queries" + }, null, 2), + }, + ], + }; + } + + // Try to find the query by IB's queryId first, then by name as fallback + let query = await this.context.flexQueryStorage.getQueryByQueryId(input.queryId); + + if (!query) { + // Try to find by name as fallback (in case user provides a friendly name) + query = await this.context.flexQueryStorage.getQueryByName(input.queryId); + } + + if (!query) { + return { + content: [ + { + type: "text", + text: JSON.stringify({ + error: "Query not found", + message: `No saved query found with ID: ${input.queryId}`, + suggestion: "Use list_flex_queries to see all saved queries" + }, null, 2), + }, + ], + }; + } + + const deleted = await this.context.flexQueryStorage.deleteQuery(query.id); + + return { + content: [ + { + type: "text", + text: JSON.stringify({ + success: deleted, + message: deleted + ? `Query "${query.name}" (${query.queryId}) has been forgotten` + : "Failed to delete query", + queryId: input.queryId, + }, null, 2), + }, + ], + }; + } catch (error) { + return { + content: [ + { + type: "text", + text: this.formatError(error), + }, + ], + }; + } + } } \ No newline at end of file diff --git a/src/tools.ts b/src/tools.ts index 9a2eedf..56d890f 100644 --- a/src/tools.ts +++ b/src/tools.ts @@ -14,7 +14,10 @@ import { GetAlertsZodShape, CreateAlertZodShape, ActivateAlertZodShape, - DeleteAlertZodShape + DeleteAlertZodShape, + GetFlexQueryZodShape, + ListFlexQueriesZodShape, + ForgetFlexQueryZodShape } from "./tool-definitions.js"; export function registerTools( @@ -135,4 +138,30 @@ export function registerTools( DeleteAlertZodShape, async (args) => await handlers.deleteAlert(args) ); + + // Register Flex Query tools (only if token is configured) + if (userConfig?.IB_FLEX_TOKEN) { + server.tool( + "get_flex_query", + "Execute a Flex Query and retrieve statements/data. The query will be automatically remembered for future use. " + + "Usage: `{ \"queryId\": \"123456\" }` or with a friendly name: `{ \"queryId\": \"123456\", \"queryName\": \"Monthly Trades\" }`. " + + "Set `parseXml: false` to get raw XML instead of parsed JSON.", + GetFlexQueryZodShape, + async (args) => await handlers.getFlexQuery(args) + ); + + server.tool( + "list_flex_queries", + "List all previously used Flex Queries that have been automatically saved. Usage: `{ \"confirm\": true }`.", + ListFlexQueriesZodShape, + async (args) => await handlers.listFlexQueries(args) + ); + + server.tool( + "forget_flex_query", + "Remove a saved Flex Query from memory. Usage: `{ \"queryId\": \"123456\" }`.", + ForgetFlexQueryZodShape, + async (args) => await handlers.forgetFlexQuery(args) + ); + } } \ No newline at end of file diff --git a/test/flex-query-client.test.ts b/test/flex-query-client.test.ts new file mode 100644 index 0000000..10dc60f --- /dev/null +++ b/test/flex-query-client.test.ts @@ -0,0 +1,378 @@ +// test/flex-query-client.test.ts +import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { FlexQueryClient } from '../src/flex-query-client.js'; +import axios from 'axios'; + +// Mock axios +vi.mock('axios'); + +describe('FlexQueryClient', () => { + let client: FlexQueryClient; + const mockToken = 'test-token-123'; + const mockQueryId = '123456'; + const mockReferenceCode = 'ref-code-789'; + + beforeEach(() => { + vi.clearAllMocks(); + client = new FlexQueryClient({ token: mockToken }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('constructor', () => { + it('should initialize with token', () => { + expect(client).toBeDefined(); + }); + + it('should create axios client with timeout', () => { + expect(axios.create).toHaveBeenCalledWith({ + timeout: 60000, + }); + }); + }); + + describe('sendRequest', () => { + it('should successfully send request and return reference code', async () => { + const mockResponse = { + data: ` + + Success + ${mockReferenceCode} + https://example.com/statement +`, + }; + + const mockGet = vi.fn().mockResolvedValue(mockResponse); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.sendRequest(mockQueryId); + + expect(mockGet).toHaveBeenCalledWith( + 'https://gdcdyn.interactivebrokers.com/Universal/servlet/FlexStatementService.SendRequest', + { + params: { + t: mockToken, + q: mockQueryId, + v: '3', + }, + } + ); + + expect(result).toEqual({ + referenceCode: mockReferenceCode, + url: 'https://example.com/statement', + }); + }); + + it('should return error on failed request', async () => { + const mockResponse = { + data: ` + + Fail + 1001 + Invalid token +`, + }; + + const mockGet = vi.fn().mockResolvedValue(mockResponse); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.sendRequest(mockQueryId); + + expect(result).toEqual({ + error: 'Invalid token', + errorCode: '1001', + }); + }); + + it('should handle network errors', async () => { + const mockError = new Error('Network error'); + Object.defineProperty(mockError, 'isAxiosError', { value: true }); + + const mockGet = vi.fn().mockRejectedValue(mockError); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + // Check the axios.isAxiosError before the test + vi.spyOn(axios, 'isAxiosError').mockReturnValue(true); + + await expect(client.sendRequest(mockQueryId)).rejects.toThrow( + 'Failed to send flex query request: Network error' + ); + }); + + it('should handle unexpected response format', async () => { + const mockResponse = { + data: ` + + Something +`, + }; + + const mockGet = vi.fn().mockResolvedValue(mockResponse); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + await expect(client.sendRequest(mockQueryId)).rejects.toThrow( + 'Unexpected response format from Flex Query service' + ); + }); + }); + + describe('getStatement', () => { + it('should successfully get statement data', async () => { + const mockStatementData = ` + + + + + + + + +`; + + const mockResponse = { + data: mockStatementData, + }; + + const mockGet = vi.fn().mockResolvedValue(mockResponse); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.getStatement(mockReferenceCode); + + expect(mockGet).toHaveBeenCalledWith( + 'https://gdcdyn.interactivebrokers.com/Universal/servlet/FlexStatementService.GetStatement', + { + params: { + t: mockToken, + q: mockReferenceCode, + v: '3', + }, + } + ); + + expect(result).toEqual({ + data: mockStatementData, + }); + }); + + it('should return error when statement not ready', async () => { + const mockResponse = { + data: ` + + Fail + 1019 + Statement generation in progress. Please try again shortly. +`, + }; + + const mockGet = vi.fn().mockResolvedValue(mockResponse); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.getStatement(mockReferenceCode); + + expect(result).toEqual({ + error: 'Statement generation in progress. Please try again shortly.', + errorCode: '1019', + }); + }); + + it('should handle network errors', async () => { + const mockError = new Error('Network timeout'); + Object.defineProperty(mockError, 'isAxiosError', { value: true }); + + const mockGet = vi.fn().mockRejectedValue(mockError); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + vi.spyOn(axios, 'isAxiosError').mockReturnValue(true); + + await expect(client.getStatement(mockReferenceCode)).rejects.toThrow( + 'Failed to get flex statement: Network timeout' + ); + }); + }); + + describe('executeQuery', () => { + it('should execute query and wait for statement', async () => { + const sendResponseXml = ` + + Success + ${mockReferenceCode} + https://example.com/statement +`; + + const statementXml = ` + + + + + +`; + + const mockGet = vi.fn() + .mockResolvedValueOnce({ data: sendResponseXml }) + .mockResolvedValueOnce({ data: statementXml }); + + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.executeQuery(mockQueryId, 3, 100); + + expect(result).toEqual({ + data: statementXml, + }); + expect(mockGet).toHaveBeenCalledTimes(2); + }); + + it('should retry when statement is not ready', async () => { + const sendResponseXml = ` + + Success + ${mockReferenceCode} +`; + + const notReadyXml = ` + + Fail + 1019 + Statement generation in progress +`; + + const statementXml = ` + + + + +`; + + const mockGet = vi.fn() + .mockResolvedValueOnce({ data: sendResponseXml }) + .mockResolvedValueOnce({ data: notReadyXml }) + .mockResolvedValueOnce({ data: statementXml }); + + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.executeQuery(mockQueryId, 3, 100); + + expect(result).toEqual({ + data: statementXml, + }); + expect(mockGet).toHaveBeenCalledTimes(3); + }); + + it('should return error when sendRequest fails', async () => { + const sendResponseXml = ` + + Fail + 1001 + Invalid query ID +`; + + const mockGet = vi.fn().mockResolvedValueOnce({ data: sendResponseXml }); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.executeQuery(mockQueryId); + + expect(result).toEqual({ + error: 'Invalid query ID', + errorCode: '1001', + }); + }); + + it('should timeout after max retries', async () => { + const sendResponseXml = ` + + Success + ${mockReferenceCode} +`; + + const notReadyXml = ` + + Fail + 1019 + Statement generation in progress +`; + + const mockGet = vi.fn() + .mockResolvedValueOnce({ data: sendResponseXml }) + .mockResolvedValue({ data: notReadyXml }); // Always return not ready + + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.executeQuery(mockQueryId, 2, 100); + + expect(result).toEqual({ + error: 'Statement not ready after 2 retries. Please try again later.', + }); + }); + + it('should return error when no reference code received', async () => { + const sendResponseXml = ` + + Success +`; + + const mockGet = vi.fn().mockResolvedValueOnce({ data: sendResponseXml }); + (axios.create as any).mockReturnValue({ get: mockGet }); + client = new FlexQueryClient({ token: mockToken }); + + const result = await client.executeQuery(mockQueryId); + + expect(result).toEqual({ + error: 'No reference code received from flex query service', + }); + }); + }); + + describe('parseStatement', () => { + it('should parse XML statement into JSON', async () => { + const xmlData = ` + + + + + + +`; + + const result = await client.parseStatement(xmlData); + + expect(result).toBeDefined(); + expect(result.FlexQueryResponse).toBeDefined(); + expect(result.FlexQueryResponse.queryName).toBe('Test Query'); + expect(result.FlexQueryResponse.type).toBe('AF'); + }); + + it('should handle parsing errors', async () => { + const invalidXml = 'This is not valid XML'; + + await expect(client.parseStatement(invalidXml)).rejects.toThrow( + 'Failed to parse flex statement XML' + ); + }); + + it('should parse XML with mergeAttrs option', async () => { + const xmlData = ` + + +`; + + const result = await client.parseStatement(xmlData); + + // mergeAttrs should merge attributes into the object + expect(result.FlexQueryResponse.Data.value).toBe('123'); + }); + }); +}); + diff --git a/test/flex-query-storage.test.ts b/test/flex-query-storage.test.ts new file mode 100644 index 0000000..2420951 --- /dev/null +++ b/test/flex-query-storage.test.ts @@ -0,0 +1,441 @@ +// test/flex-query-storage.test.ts +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import fs from 'fs/promises'; +import os from 'os'; +import path from 'path'; + +// Mock fs and os BEFORE importing the module +vi.mock('fs/promises'); +vi.mock('os', () => ({ + default: { + homedir: vi.fn(() => '/mock/home') + }, + homedir: vi.fn(() => '/mock/home') +})); + +import { FlexQueryStorage, SavedFlexQuery } from '../src/flex-query-storage.js'; + +describe('FlexQueryStorage', () => { + let storage: FlexQueryStorage; + const mockStorageFile = '/mock/path/flex-queries.json'; + const mockHomeDir = '/mock/home'; + + beforeEach(() => { + vi.clearAllMocks(); + + // Mock os.homedir + (os.homedir as any).mockReturnValue(mockHomeDir); + + // Mock fs operations + (fs.access as any).mockResolvedValue(undefined); + (fs.mkdir as any).mockResolvedValue(undefined); + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + (fs.writeFile as any).mockResolvedValue(undefined); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('constructor', () => { + it('should use default storage file path', () => { + storage = new FlexQueryStorage(); + const expectedPath = path.join(mockHomeDir, '.interactive-brokers-mcp', 'flex-queries.json'); + expect(storage.getStorageFilePath()).toBe(expectedPath); + }); + + it('should use custom storage file path', () => { + storage = new FlexQueryStorage(mockStorageFile); + expect(storage.getStorageFilePath()).toBe(mockStorageFile); + }); + }); + + describe('initialize', () => { + it('should create directory when it does not exist', async () => { + // Setup mocks before initialization + (fs.access as any).mockRejectedValueOnce(new Error('ENOENT')); // Directory doesn't exist + (fs.readFile as any).mockRejectedValueOnce({ code: 'ENOENT' }); // File doesn't exist + + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + + // Should create the directory + expect(fs.mkdir).toHaveBeenCalledWith( + path.dirname(mockStorageFile), + { recursive: true } + ); + + // Should NOT write file yet (file is created lazily on first save) + expect(fs.writeFile).not.toHaveBeenCalled(); + }); + + it('should load existing data if file exists', async () => { + const existingData = { + queries: [ + { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }, + ], + }; + + // Mock fs.access to succeed (directory exists) + (fs.access as any).mockResolvedValue(undefined); + // Mock readFile to return existing data (called twice: once for initialize, once for listQueries) + (fs.readFile as any).mockResolvedValue(JSON.stringify(existingData)); + + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + + const queries = await storage.listQueries(); + expect(queries).toHaveLength(1); + expect(queries[0].name).toBe('Test Query'); + }); + + it('should handle initialization errors', async () => { + (fs.access as any).mockRejectedValue(new Error('Permission denied')); + (fs.mkdir as any).mockRejectedValue(new Error('Cannot create directory')); + + storage = new FlexQueryStorage(mockStorageFile); + + await expect(storage.initialize()).rejects.toThrow(); + }); + }); + + describe('listQueries', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should return empty array when no queries', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const queries = await storage.listQueries(); + expect(queries).toEqual([]); + }); + + it('should return all saved queries', async () => { + const mockQueries = [ + { + id: 'query_1', + name: 'Query 1', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }, + { + id: 'query_2', + name: 'Query 2', + queryId: '456', + createdAt: '2023-01-02T00:00:00.000Z', + }, + ]; + + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: mockQueries })); + + const queries = await storage.listQueries(); + expect(queries).toHaveLength(2); + expect(queries[0].name).toBe('Query 1'); + expect(queries[1].name).toBe('Query 2'); + }); + }); + + describe('getQuery', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should get query by internal ID', async () => { + const mockQueries = [ + { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }, + ]; + + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: mockQueries })); + + const query = await storage.getQuery('query_1'); + expect(query).toBeDefined(); + expect(query?.name).toBe('Test Query'); + }); + + it('should return null when query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const query = await storage.getQuery('nonexistent'); + expect(query).toBeNull(); + }); + }); + + describe('getQueryByQueryId', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should get query by IB query ID', async () => { + const mockQueries = [ + { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }, + ]; + + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: mockQueries })); + + const query = await storage.getQueryByQueryId('123456'); + expect(query).toBeDefined(); + expect(query?.name).toBe('Test Query'); + }); + + it('should return null when query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const query = await storage.getQueryByQueryId('nonexistent'); + expect(query).toBeNull(); + }); + }); + + describe('getQueryByName', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should get query by name (case insensitive)', async () => { + const mockQueries = [ + { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }, + ]; + + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: mockQueries })); + + const query = await storage.getQueryByName('test query'); + expect(query).toBeDefined(); + expect(query?.name).toBe('Test Query'); + }); + + it('should return null when query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const query = await storage.getQueryByName('nonexistent'); + expect(query).toBeNull(); + }); + }); + + describe('saveQuery', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should save a new query', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const newQuery = { + name: 'New Query', + queryId: '789', + description: 'Test description', + }; + + const savedQuery = await storage.saveQuery(newQuery); + + expect(savedQuery).toBeDefined(); + expect(savedQuery.id).toBeDefined(); + expect(savedQuery.name).toBe('New Query'); + expect(savedQuery.queryId).toBe('789'); + expect(savedQuery.createdAt).toBeDefined(); + expect(fs.writeFile).toHaveBeenCalled(); + }); + + it('should throw error when queryId already exists', async () => { + const existingQuery = { + id: 'query_1', + name: 'Existing Query', + queryId: '789', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + (fs.readFile as any).mockResolvedValue( + JSON.stringify({ queries: [existingQuery] }) + ); + + const newQuery = { + name: 'New Query', + queryId: '789', // Same as existing + description: 'Test description', + }; + + await expect(storage.saveQuery(newQuery)).rejects.toThrow( + 'A query with the IB Query ID "789" already exists' + ); + }); + + it('should generate unique ID for each query', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const query1 = await storage.saveQuery({ + name: 'Query 1', + queryId: '123', + }); + + const query2 = await storage.saveQuery({ + name: 'Query 2', + queryId: '456', + }); + + expect(query1.id).not.toBe(query2.id); + }); + }); + + describe('updateQuery', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should update query fields', async () => { + const existingQuery = { + id: 'query_1', + name: 'Old Name', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + (fs.readFile as any).mockResolvedValue( + JSON.stringify({ queries: [existingQuery] }) + ); + + const updated = await storage.updateQuery('query_1', { + name: 'New Name', + description: 'Updated description', + }); + + expect(updated.name).toBe('New Name'); + expect(updated.description).toBe('Updated description'); + expect(updated.id).toBe('query_1'); + expect(fs.writeFile).toHaveBeenCalled(); + }); + + it('should throw error when query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + await expect( + storage.updateQuery('nonexistent', { name: 'New Name' }) + ).rejects.toThrow('Query with ID "nonexistent" not found'); + }); + + it('should prevent duplicate names', async () => { + const queries = [ + { + id: 'query_1', + name: 'Query 1', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }, + { + id: 'query_2', + name: 'Query 2', + queryId: '456', + createdAt: '2023-01-02T00:00:00.000Z', + }, + ]; + + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries })); + + await expect( + storage.updateQuery('query_1', { name: 'Query 2' }) + ).rejects.toThrow('A query with the name "Query 2" already exists'); + }); + }); + + describe('markQueryUsed', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should update lastUsed timestamp', async () => { + const existingQuery = { + id: 'query_1', + name: 'Test Query', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + (fs.readFile as any).mockResolvedValue( + JSON.stringify({ queries: [existingQuery] }) + ); + + await storage.markQueryUsed('query_1'); + + expect(fs.writeFile).toHaveBeenCalled(); + // Check that the query now has a lastUsed field + const savedData = JSON.parse((fs.writeFile as any).mock.calls[0][1]); + expect(savedData.queries[0].lastUsed).toBeDefined(); + }); + + it('should not throw error if query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + await expect(storage.markQueryUsed('nonexistent')).resolves.not.toThrow(); + }); + }); + + describe('deleteQuery', () => { + beforeEach(async () => { + storage = new FlexQueryStorage(mockStorageFile); + await storage.initialize(); + }); + + it('should delete query and return true', async () => { + const existingQuery = { + id: 'query_1', + name: 'Test Query', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + (fs.readFile as any).mockResolvedValue( + JSON.stringify({ queries: [existingQuery] }) + ); + + const result = await storage.deleteQuery('query_1'); + + expect(result).toBe(true); + expect(fs.writeFile).toHaveBeenCalled(); + + // Verify query was removed + const savedData = JSON.parse((fs.writeFile as any).mock.calls[0][1]); + expect(savedData.queries).toHaveLength(0); + }); + + it('should return false when query not found', async () => { + (fs.readFile as any).mockResolvedValue(JSON.stringify({ queries: [] })); + + const result = await storage.deleteQuery('nonexistent'); + + expect(result).toBe(false); + }); + }); + + describe('getStorageFilePath', () => { + it('should return the storage file path', () => { + storage = new FlexQueryStorage(mockStorageFile); + expect(storage.getStorageFilePath()).toBe(mockStorageFile); + }); + }); +}); + diff --git a/test/tool-handlers.test.ts b/test/tool-handlers.test.ts index ccbab64..9d42022 100644 --- a/test/tool-handlers.test.ts +++ b/test/tool-handlers.test.ts @@ -280,5 +280,311 @@ describe('ToolHandlers', () => { expect(result.content[0].text).toContain('String error'); }); }); + + describe('Flex Query Tools', () => { + describe('getFlexQuery', () => { + it('should return error when flex query client is not configured', async () => { + // Context without flex query client (using the one from beforeEach which has no flex client) + const result = await handlers.getFlexQuery({ + queryId: '123456', + parseXml: false, + }); + + const response = JSON.parse(result.content[0].text); + expect(response.error).toBe('Flex Query feature not configured'); + expect(response.message).toContain('IB_FLEX_TOKEN'); + }); + + it('should execute flex query when configured', async () => { + // Create a fresh context with flex query configuration + const mockFlexQueryClient = { + executeQuery: vi.fn().mockResolvedValue({ + data: '', + }), + parseStatement: vi.fn().mockResolvedValue({ + FlexQueryResponse: { + queryName: 'Test Query', + FlexStatements: {}, + }, + }), + }; + + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(null), + saveQuery: vi.fn().mockResolvedValue({ + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }), + markQueryUsed: vi.fn().mockResolvedValue(undefined), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + // Create NEW context with flex query setup + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryClient: mockFlexQueryClient as any, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.getFlexQuery({ + queryId: '123456', + parseXml: false, + }); + + const response = JSON.parse(result.content[0].text); + expect(response.success).toBe(true); + expect(response.queryId).toBe('123456'); + expect(response.autoSaved).toBe(true); + expect(mockFlexQueryClient.executeQuery).toHaveBeenCalledWith('123456'); + expect(mockFlexQueryStorage.saveQuery).toHaveBeenCalled(); + }); + + it('should mark query as used when already exists', async () => { + const existingQuery = { + id: 'query_1', + name: 'Existing Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + const mockFlexQueryClient = { + executeQuery: vi.fn().mockResolvedValue({ + data: '', + }), + parseStatement: vi.fn().mockResolvedValue({ + FlexQueryResponse: { + queryName: 'Test Query', + }, + }), + }; + + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(existingQuery), + markQueryUsed: vi.fn().mockResolvedValue(undefined), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryClient: mockFlexQueryClient as any, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.getFlexQuery({ + queryId: '123456', + parseXml: false, + }); + + const response = JSON.parse(result.content[0].text); + expect(response.autoSaved).toBe(false); + expect(mockFlexQueryStorage.markQueryUsed).toHaveBeenCalledWith('query_1'); + }); + + it('should handle flex query errors', async () => { + const mockFlexQueryClient = { + executeQuery: vi.fn().mockResolvedValue({ + error: 'Invalid query ID', + errorCode: '1001', + }), + }; + + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(null), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryClient: mockFlexQueryClient as any, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.getFlexQuery({ + queryId: '123456', + parseXml: false, + }); + + const response = JSON.parse(result.content[0].text); + expect(response.error).toBe('Invalid query ID'); + expect(response.errorCode).toBe('1001'); + }); + }); + + describe('listFlexQueries', () => { + it('should return error when not configured', async () => { + const result = await handlers.listFlexQueries({ confirm: true }); + + const response = JSON.parse(result.content[0].text); + expect(response.error).toBe('Flex Query feature not configured'); + }); + + it('should list all saved queries', async () => { + const mockQueries = [ + { + id: 'query_1', + name: 'Query 1', + queryId: '123', + createdAt: '2023-01-01T00:00:00.000Z', + }, + { + id: 'query_2', + name: 'Query 2', + queryId: '456', + createdAt: '2023-01-02T00:00:00.000Z', + }, + ]; + + const mockFlexQueryStorage = { + listQueries: vi.fn().mockResolvedValue(mockQueries), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path/flex-queries.json'), + }; + + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.listFlexQueries({ confirm: true }); + + const response = JSON.parse(result.content[0].text); + expect(response.count).toBe(2); + expect(response.queries).toHaveLength(2); + expect(response.storageLocation).toBe('/mock/path/flex-queries.json'); + }); + }); + + describe('forgetFlexQuery', () => { + it('should return error when not configured', async () => { + const result = await handlers.forgetFlexQuery({ queryId: '123456' }); + + const response = JSON.parse(result.content[0].text); + expect(response.error).toBe('Flex Query feature not configured'); + }); + + it('should delete query by queryId', async () => { + const existingQuery = { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(existingQuery), + getQueryByName: vi.fn().mockResolvedValue(null), + deleteQuery: vi.fn().mockResolvedValue(true), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.forgetFlexQuery({ queryId: '123456' }); + + const response = JSON.parse(result.content[0].text); + expect(response.success).toBe(true); + expect(response.message).toContain('Test Query'); + expect(mockFlexQueryStorage.deleteQuery).toHaveBeenCalledWith('query_1'); + }); + + it('should try name lookup as fallback', async () => { + const existingQuery = { + id: 'query_1', + name: 'Test Query', + queryId: '123456', + createdAt: '2023-01-01T00:00:00.000Z', + }; + + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(null), + getQueryByName: vi.fn().mockResolvedValue(existingQuery), + deleteQuery: vi.fn().mockResolvedValue(true), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + const flexContext: ToolHandlerContext = { + ibClient: mockIBClient, + gatewayManager: mockGatewayManager, + config: { + ...context.config, + IB_FLEX_TOKEN: 'test-token', + }, + flexQueryStorage: mockFlexQueryStorage as any, + }; + + const flexHandlers = new ToolHandlers(flexContext); + + const result = await flexHandlers.forgetFlexQuery({ queryId: 'Test Query' }); + + const response = JSON.parse(result.content[0].text); + expect(response.success).toBe(true); + expect(mockFlexQueryStorage.getQueryByName).toHaveBeenCalledWith('Test Query'); + }); + + it('should return error when query not found', async () => { + const mockFlexQueryStorage = { + getQueryByQueryId: vi.fn().mockResolvedValue(null), + getQueryByName: vi.fn().mockResolvedValue(null), + initialize: vi.fn().mockResolvedValue(undefined), + getStorageFilePath: vi.fn().mockReturnValue('/mock/path'), + }; + + context.config.IB_FLEX_TOKEN = 'test-token'; + context.flexQueryStorage = mockFlexQueryStorage as any; + handlers = new ToolHandlers(context); + + const result = await handlers.forgetFlexQuery({ queryId: 'nonexistent' }); + + const response = JSON.parse(result.content[0].text); + expect(response.error).toBe('Query not found'); + expect(response.message).toContain('nonexistent'); + }); + }); + }); });