diff --git a/src/dws/credits.ts b/src/dws/credits.ts new file mode 100644 index 0000000..87ff6b4 --- /dev/null +++ b/src/dws/credits.ts @@ -0,0 +1,76 @@ +import axios from 'axios' +import { getApiKey, pipeToString } from './utils.js' +import { getVersion } from '../version.js' +import { CallToolResult } from '@modelcontextprotocol/sdk/types.js' + +/** + * Account info response from DWS API (GET /account/info) + * + * Reference: https://www.nutrient.io/api/reference/public/#tag/Account/operation/get-account-info + */ +interface AccountInfoResponse { + apiKeys?: { + live?: string + } + signedIn?: boolean + subscriptionType?: string + usage?: { + totalCredits?: number + usedCredits?: number + } +} + +/** + * Strips sensitive fields (API keys) from the account info response + * before returning it to the LLM. + */ +export function sanitizeAccountInfo(data: AccountInfoResponse): Omit { + const { apiKeys: _apiKeys, ...safe } = data + return safe +} + +/** + * Calls the DWS /account/info endpoint and returns credit information. + */ +export async function performCheckCreditsCall(): Promise { + const apiKey = getApiKey() + + const response = await axios.get('https://api.nutrient.io/account/info', { + headers: { + Authorization: `Bearer ${apiKey}`, + 'User-Agent': `NutrientDWSMCPServer/${getVersion()}`, + }, + responseType: 'stream', + }) + + const raw = await pipeToString(response.data) + + let parsed: AccountInfoResponse + try { + parsed = JSON.parse(raw) + } catch { + return { + content: [{ type: 'text', text: `Unexpected non-JSON response from /account/info: ${raw}` }], + isError: true, + } + } + + const safe = sanitizeAccountInfo(parsed) + + const totalCredits = safe.usage?.totalCredits + const usedCredits = safe.usage?.usedCredits + const remainingCredits = totalCredits != null && usedCredits != null ? totalCredits - usedCredits : undefined + + const summary = { + subscriptionType: safe.subscriptionType ?? 'unknown', + totalCredits: totalCredits ?? 'unknown', + usedCredits: usedCredits ?? 'unknown', + remainingCredits: remainingCredits ?? 'unknown', + signedIn: safe.signedIn, + } + + return { + content: [{ type: 'text', text: JSON.stringify(summary, null, 2) }], + isError: false, + } +} diff --git a/src/index.ts b/src/index.ts index 2ed80a3..4d7f537 100644 --- a/src/index.ts +++ b/src/index.ts @@ -8,10 +8,17 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' -import { AiRedactArgsSchema, BuildAPIArgsSchema, DirectoryTreeArgsSchema, SignAPIArgsSchema } from './schemas.js' +import { + AiRedactArgsSchema, + BuildAPIArgsSchema, + CheckCreditsArgsSchema, + DirectoryTreeArgsSchema, + SignAPIArgsSchema, +} from './schemas.js' import { performBuildCall } from './dws/build.js' import { performSignCall } from './dws/sign.js' import { performAiRedactCall } from './dws/ai-redact.js' +import { performCheckCreditsCall } from './dws/credits.js' import { performDirectoryTreeCall } from './fs/directoryTree.js' import { setSandboxDirectory } from './fs/sandbox.js' import { createErrorResponse } from './responses.js' @@ -105,6 +112,21 @@ By default (when neither stage nor apply is set), redactions are detected and im }, ) + server.tool( + 'check_credits', + `Check your Nutrient DWS API credit balance and usage for the current billing period. + +Returns: subscription type, total credits, used credits, and remaining credits.`, + CheckCreditsArgsSchema.shape, + async () => { + try { + return performCheckCreditsCall() + } catch (error) { + return createErrorResponse(`Error: ${error instanceof Error ? error.message : String(error)}`) + } + }, + ) + if (sandboxEnabled) { server.tool( 'sandbox_file_tree', diff --git a/src/schemas.ts b/src/schemas.ts index 9e97706..dbd43a3 100644 --- a/src/schemas.ts +++ b/src/schemas.ts @@ -1,5 +1,9 @@ import { z } from 'zod' +export const CheckCreditsArgsSchema = z.object({}) + +export type CheckCreditsArgs = z.infer + export const DirectoryTreeArgsSchema = z.object({ path: z .string() diff --git a/tests/unit.test.ts b/tests/unit.test.ts index 6f1773e..b924728 100644 --- a/tests/unit.test.ts +++ b/tests/unit.test.ts @@ -1,7 +1,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import fs, { Stats } from 'fs' import { Readable } from 'stream' -import { AiRedactArgsSchema, Instructions, SignatureOptions } from '../src/schemas.js' +import { AiRedactArgsSchema, CheckCreditsArgsSchema, Instructions, SignatureOptions } from '../src/schemas.js' +import { sanitizeAccountInfo, performCheckCreditsCall } from '../src/dws/credits.js' import { config as dotenvConfig } from 'dotenv' import { performBuildCall } from '../src/dws/build.js' import { performSignCall } from '../src/dws/sign.js' @@ -820,4 +821,114 @@ describe('API Functions', () => { expect(result).toBeUndefined() }) }) + + describe('CheckCreditsArgsSchema', () => { + it('should accept an empty object', () => { + const result = CheckCreditsArgsSchema.safeParse({}) + expect(result.success).toBe(true) + }) + }) + + describe('sanitizeAccountInfo', () => { + it('should strip apiKeys from the response', () => { + const response = { + apiKeys: { live: 'sk_live_SUPER_SECRET_KEY_12345' }, + signedIn: true, + subscriptionType: 'paid', + usage: { totalCredits: 100, usedCredits: 42 }, + } + + const sanitized = sanitizeAccountInfo(response) + + expect(sanitized).not.toHaveProperty('apiKeys') + expect(sanitized).toEqual({ + signedIn: true, + subscriptionType: 'paid', + usage: { totalCredits: 100, usedCredits: 42 }, + }) + }) + + it('should handle response with no apiKeys field', () => { + const response = { + signedIn: true, + subscriptionType: 'free', + usage: { totalCredits: 100, usedCredits: 0 }, + } + + const sanitized = sanitizeAccountInfo(response) + + expect(sanitized).not.toHaveProperty('apiKeys') + expect(sanitized.subscriptionType).toBe('free') + }) + + it('should never leak the live API key to the LLM', () => { + const secretKey = 'sk_live_MUST_NEVER_APPEAR_IN_OUTPUT' + const response = { + apiKeys: { live: secretKey }, + signedIn: true, + subscriptionType: 'enterprise', + usage: { totalCredits: 10000, usedCredits: 500 }, + } + + const sanitized = sanitizeAccountInfo(response) + const serialized = JSON.stringify(sanitized) + + expect(serialized).not.toContain(secretKey) + expect(serialized).not.toContain('apiKeys') + expect(serialized).not.toContain('sk_live') + }) + }) + + describe('performCheckCreditsCall', () => { + it('should return credit summary from API response', async () => { + const apiResponse = { + apiKeys: { live: 'sk_live_secret' }, + signedIn: true, + subscriptionType: 'paid', + usage: { totalCredits: 100, usedCredits: 42 }, + } + + vi.stubEnv('NUTRIENT_DWS_API_KEY', 'test-key') + vi.spyOn(axios, 'get').mockResolvedValue({ + data: Readable.from([JSON.stringify(apiResponse)]), + }) + + const result = await performCheckCreditsCall() + + expect(result.isError).toBe(false) + const text = (result.content[0] as TextContent).text + const parsed = JSON.parse(text) + expect(parsed.subscriptionType).toBe('paid') + expect(parsed.totalCredits).toBe(100) + expect(parsed.usedCredits).toBe(42) + expect(parsed.remainingCredits).toBe(58) + // Must not contain the API key + expect(text).not.toContain('sk_live_secret') + + vi.restoreAllMocks() + }) + + it('should handle non-JSON API response', async () => { + vi.stubEnv('NUTRIENT_DWS_API_KEY', 'test-key') + vi.spyOn(axios, 'get').mockResolvedValue({ + data: Readable.from(['not json']), + }) + + const result = await performCheckCreditsCall() + + expect(result.isError).toBe(true) + expect((result.content[0] as TextContent).text).toContain('Unexpected non-JSON response') + + vi.restoreAllMocks() + }) + + it('should error when API key is not set', async () => { + vi.stubEnv('NUTRIENT_DWS_API_KEY', '') + delete process.env.NUTRIENT_DWS_API_KEY + + await expect(performCheckCreditsCall()).rejects.toThrow('NUTRIENT_DWS_API_KEY not set') + + vi.restoreAllMocks() + }) + }) })