feat: add check_credits tool via /account/info API (v2)

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<AccountInfoResponse, 'apiKeys'> {
+  const { apiKeys: _apiKeys, ...safe } = data
+  return safe
+}
+
+/**
+ * Calls the DWS /account/info endpoint and returns credit information.
+ */
+export async function performCheckCreditsCall(): Promise<CallToolResult> {
+  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,
+  }
+}

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',

src/schemas.ts

@@ -1,5 +1,9 @@
 import { z } from 'zod'
 
+export const CheckCreditsArgsSchema = z.object({})
+
+export type CheckCreditsArgs = z.infer<typeof CheckCreditsArgsSchema>
+
 export const DirectoryTreeArgsSchema = z.object({
   path: z
     .string()

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()
+    })
+  })
 })