feat: add credit tracking via API (check_credits tool)

src/dws/api.ts

@@ -7,25 +7,36 @@ import { getVersion } from '../version.js'
  * Makes an API call to the Nutrient API
  * @param endpoint The API endpoint to call (e.g., 'sign', 'build')
  * @param data The data to send (FormData or JSON object)
+ * @param options Optional request settings (e.g. HTTP method)
  * @returns The API response
  */
-export async function callNutrientApi(endpoint: string, data: FormData | Record<string, unknown>) {
+export async function callNutrientApi(
+  endpoint: string,
+  data?: FormData | Record<string, unknown>,
+  options: { method?: 'GET' | 'POST' } = {},
+) {
   const apiKey = getApiKey()
   const isFormData = data instanceof FormData
+  const method = options.method ?? 'POST'
 
   const defaultHeaders: Record<string, string> = {
     Authorization: `Bearer ${apiKey}`,
     'User-Agent': `NutrientDWSMCPServer/${getVersion()}`,
   }
 
-  const headers: Record<string, string> = isFormData
-    ? defaultHeaders
-    : {
-        ...defaultHeaders,
-        'Content-Type': 'application/json',
-      }
-  return axios.post(`https://api.nutrient.io/${endpoint}`, data, {
+  const headers: Record<string, string> =
+    isFormData || method === 'GET'
+      ? defaultHeaders
+      : {
+          ...defaultHeaders,
+          'Content-Type': 'application/json',
+        }
+
+  return axios.request({
+    method,
+    url: `https://api.nutrient.io/${endpoint}`,
     headers,
     responseType: 'stream',
+    ...(method === 'GET' ? {} : { data: data ?? {} }),
   })
 }

src/dws/credits.ts

@@ -0,0 +1,135 @@
+import { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import { callNutrientApi } from './api.js'
+import { handleApiError, pipeToString } from './utils.js'
+import { createSuccessResponse } from '../responses.js'
+
+export type CreditAction = 'balance' | 'usage'
+export type CreditPeriod = 'day' | 'week' | 'month' | 'all'
+
+const DEFAULT_USAGE_PERIOD: CreditPeriod = 'week'
+
+const BALANCE_PATHS: string[][] = [
+  ['credits', 'remaining'],
+  ['credits', 'balance'],
+  ['credits', 'remainingCredits'],
+  ['credits', 'remaining_credits'],
+  ['remainingCredits'],
+  ['remaining_credits'],
+  ['remaining'],
+  ['balance'],
+  ['creditBalance'],
+  ['credit_balance'],
+]
+
+const USAGE_PATHS: string[][] = [
+  ['usage'],
+  ['credits', 'usage'],
+  ['consumption'],
+  ['creditUsage'],
+  ['usageBreakdown'],
+  ['usage_breakdown'],
+]
+
+/**
+ * Fetches current account credit information from the Nutrient API.
+ */
+export async function performCheckCreditsCall(
+  action: CreditAction,
+  period?: CreditPeriod,
+): Promise<CallToolResult> {
+  try {
+    const response = await callNutrientApi('account/info', undefined, { method: 'GET' })
+    const responseString = await pipeToString(response.data)
+
+    let payload: unknown
+    try {
+      payload = JSON.parse(responseString)
+    } catch {
+      return createSuccessResponse(responseString)
+    }
+
+    if (action === 'balance') {
+      const balance = extractBalance(payload)
+      if (balance !== null) {
+        return createSuccessResponse(JSON.stringify({ balance }, null, 2))
+      }
+      return createSuccessResponse(JSON.stringify(payload, null, 2))
+    }
+
+    const usagePeriod = period ?? DEFAULT_USAGE_PERIOD
+    const usage = extractUsage(payload, usagePeriod)
+    if (usage !== null) {
+      return createSuccessResponse(JSON.stringify({ period: usagePeriod, usage }, null, 2))
+    }
+
+    return createSuccessResponse(JSON.stringify(payload, null, 2))
+  } catch (e: unknown) {
+    return handleApiError(e)
+  }
+}
+
+function extractBalance(payload: unknown): number | null {
+  if (!isRecord(payload)) {
+    return null
+  }
+
+  for (const path of BALANCE_PATHS) {
+    const value = getNestedValue(payload, path)
+    const numberValue = parseNumber(value)
+    if (numberValue !== null) {
+      return numberValue
+    }
+  }
+
+  return null
+}
+
+function extractUsage(payload: unknown, period: CreditPeriod): unknown | null {
+  if (!isRecord(payload)) {
+    return null
+  }
+
+  for (const path of USAGE_PATHS) {
+    const value = getNestedValue(payload, path)
+    if (isRecord(value)) {
+      if (period in value) {
+        return value[period]
+      }
+      return { allPeriods: value, note: "Requested period not available, returning all usage data" }
+    }
+  }
+
+  return null
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === 'object' && value !== null
+}
+
+function getNestedValue(source: Record<string, unknown>, path: string[]): unknown {
+  let current: unknown = source
+
+  for (const key of path) {
+    if (!isRecord(current) || !(key in current)) {
+      return undefined
+    }
+    current = current[key]
+  }
+
+  return current
+}
+
+function parseNumber(value: unknown): number | null {
+  if (typeof value === 'number' && Number.isFinite(value)) {
+    return value
+  }
+
+  if (typeof value === 'string') {
+    const parsed = Number.parseFloat(value)
+    if (Number.isFinite(parsed)) {
+      return parsed
+    }
+  }
+
+  return null
+}

src/index.ts

@@ -8,9 +8,10 @@
 
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
-import { BuildAPIArgsSchema, DirectoryTreeArgsSchema, SignAPIArgsSchema } from './schemas.js'
+import { BuildAPIArgsSchema, CheckCreditsArgsSchema, DirectoryTreeArgsSchema, SignAPIArgsSchema } from './schemas.js'
 import { performBuildCall } from './dws/build.js'
 import { performSignCall } from './dws/sign.js'
+import { performCheckCreditsCall } from './dws/credits.js'
 import { performDirectoryTreeCall } from './fs/directoryTree.js'
 import { setSandboxDirectory } from './fs/sandbox.js'
 import { createErrorResponse } from './responses.js'
@@ -81,6 +82,23 @@ Positioning:
     },
   )
 
+  server.tool(
+    'check_credits',
+    `Checks Nutrient DWS account credits using the /account/info endpoint.
+
+Actions:
+• balance - remaining credits
+• usage - usage breakdown for a period (day, week, month, all; defaults to week if not specified)`,
+    CheckCreditsArgsSchema.shape,
+    async ({ action, period }) => {
+      try {
+        return performCheckCreditsCall(action, period)
+      } catch (error) {
+        return createErrorResponse(`Error: ${error instanceof Error ? error.message : String(error)}`)
+      }
+    },
+  )
+
   if (sandboxEnabled) {
     server.tool(
       'sandbox_file_tree',

src/schemas.ts

@@ -8,6 +8,16 @@ export const DirectoryTreeArgsSchema = z.object({
     ),
 })
 
+export const CheckCreditsArgsSchema = z.object({
+  action: z
+    .enum(['balance', 'usage'])
+    .describe('The type of credit information to retrieve: balance for remaining credits or usage for consumption.'),
+  period: z
+    .enum(['day', 'week', 'month', 'all'])
+    .optional()
+    .describe('Time period to return usage for. Defaults to week when action is usage.'),
+})
+
 export const RectSchema = z
   .array(z.number())
   .length(4)
@@ -508,3 +518,4 @@ export type Action = z.infer<typeof BuildActionSchema>
 
 export type SignAPIArgs = z.infer<typeof SignAPIArgsSchema>
 export type SignatureOptions = z.infer<typeof CreateDigitalSignatureSchema>
+export type CheckCreditsArgs = z.infer<typeof CheckCreditsArgsSchema>

tests/unit.test.ts

@@ -1,9 +1,10 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import fs, { Stats } from 'fs'
 import { Readable } from 'stream'
-import { Instructions, SignatureOptions } from '../src/schemas.js'
+import { CheckCreditsArgsSchema, Instructions, SignatureOptions } from '../src/schemas.js'
 import { config as dotenvConfig } from 'dotenv'
 import { performBuildCall } from '../src/dws/build.js'
+import { performCheckCreditsCall } from '../src/dws/credits.js'
 import { performSignCall } from '../src/dws/sign.js'
 import { performDirectoryTreeCall } from '../src/fs/directoryTree.js'
 import * as sandbox from '../src/fs/sandbox.js'
@@ -364,6 +365,87 @@ describe('API Functions', () => {
     })
   })
 
+  describe('performCheckCreditsCall', () => {
+    it('should extract balance from nested JSON', async () => {
+      const mockStream = createMockStream(JSON.stringify({ credits: { remaining: 500 } }))
+      vi.mocked(api.callNutrientApi).mockResolvedValueOnce({
+        data: mockStream,
+        status: 200,
+        statusText: 'OK',
+        headers: {},
+        config: {} as InternalAxiosRequestConfig,
+      })
+
+      const result = await performCheckCreditsCall('balance')
+
+      expect(result.isError).toBe(false)
+      expect(api.callNutrientApi).toHaveBeenCalledWith('account/info', undefined, { method: 'GET' })
+
+      const payload = JSON.parse(getTextContent(result))
+      expect(payload).toEqual({ balance: 500 })
+    })
+
+    it('should extract usage for matching period', async () => {
+      const mockStream = createMockStream(JSON.stringify({ usage: { week: { used: 10 }, month: { used: 33 } } }))
+      vi.mocked(api.callNutrientApi).mockResolvedValueOnce({
+        data: mockStream,
+        status: 200,
+        statusText: 'OK',
+        headers: {},
+        config: {} as InternalAxiosRequestConfig,
+      })
+
+      const result = await performCheckCreditsCall('usage', 'month')
+
+      expect(result.isError).toBe(false)
+
+      const payload = JSON.parse(getTextContent(result))
+      expect(payload).toEqual({ period: 'month', usage: { used: 33 } })
+    })
+
+    it('should return raw text for non-JSON responses', async () => {
+      const mockStream = createMockStream('not json')
+      vi.mocked(api.callNutrientApi).mockResolvedValueOnce({
+        data: mockStream,
+        status: 200,
+        statusText: 'OK',
+        headers: {},
+        config: {} as InternalAxiosRequestConfig,
+      })
+
+      const result = await performCheckCreditsCall('balance')
+
+      expect(result.isError).toBe(false)
+      expect(getTextContent(result)).toBe('not json')
+    })
+
+    it('should handle API errors', async () => {
+      vi.mocked(api.callNutrientApi).mockRejectedValueOnce(new Error('Boom'))
+
+      const result = await performCheckCreditsCall('balance')
+
+      expect(result.isError).toBe(true)
+      expect(getTextContent(result)).toContain('Error: Boom')
+    })
+
+    it('should return full JSON when balance cannot be extracted', async () => {
+      const apiPayload = { credits: { spent: 12 }, other: 'value' }
+      const mockStream = createMockStream(JSON.stringify(apiPayload))
+      vi.mocked(api.callNutrientApi).mockResolvedValueOnce({
+        data: mockStream,
+        status: 200,
+        statusText: 'OK',
+        headers: {},
+        config: {} as InternalAxiosRequestConfig,
+      })
+
+      const result = await performCheckCreditsCall('balance')
+
+      expect(result.isError).toBe(false)
+      expect(JSON.parse(getTextContent(result))).toEqual(apiPayload)
+    })
+  })
+
   describe('performDirectoryTreeCall', () => {
     it('should return a tree structure for a valid directory', async () => {
       vi.spyOn(fs.promises, 'access').mockImplementation(async () => {})
@@ -499,6 +581,31 @@ describe('API Functions', () => {
     })
   })
 
+  describe('CheckCreditsArgsSchema', () => {
+    it('should accept balance without period', () => {
+      const result = CheckCreditsArgsSchema.parse({ action: 'balance' })
+      expect(result.action).toBe('balance')
+      expect(result.period).toBeUndefined()
+    })
+
+    it('should accept usage with explicit period', () => {
+      const result = CheckCreditsArgsSchema.parse({ action: 'usage', period: 'month' })
+      expect(result.action).toBe('usage')
+      expect(result.period).toBe('month')
+    })
+
+    it('should allow usage without period', () => {
+      const result = CheckCreditsArgsSchema.parse({ action: 'usage' })
+      expect(result.action).toBe('usage')
+      expect(result.period).toBeUndefined()
+    })
+
+    it('should reject invalid action', () => {
+      const result = CheckCreditsArgsSchema.safeParse({ action: 'invalid' })
+      expect(result.success).toBe(false)
+    })
+  })
+
   describe('Sandbox Functionality', () => {
     beforeEach(async () => {
       vi.resetAllMocks()