diff --git a/SDK.md b/SDK.md index ed669983..e1ed3928 100644 --- a/SDK.md +++ b/SDK.md @@ -209,10 +209,24 @@ for (const item of results.organic) { ### Quota ```typescript +import { + isAccountBalanceResponse, + isQuotaResponse, +} from 'mmx-cli/sdk'; + const quota = await sdk.quota.info(); -console.log(quota); + +if (isQuotaResponse(quota)) { + console.log(quota.model_remains); +} else if (isAccountBalanceResponse(quota)) { + console.log(quota.available_amount); +} ``` +Token Plan credentials and OAuth return `QuotaResponse`. API secret keys with +the `sk-api-` prefix return `AccountBalanceResponse`. The type guards above +narrow the `QuotaInfoResponse` union without changing the raw API response. + ## Custom Base URL ```typescript diff --git a/src/client/endpoints.ts b/src/client/endpoints.ts index bde656df..8627fd03 100644 --- a/src/client/endpoints.ts +++ b/src/client/endpoints.ts @@ -62,10 +62,37 @@ export function isSecretApiKey(apiKey: string): boolean { return apiKey.startsWith('sk-api-'); } +export type UsageEndpointKind = 'quota' | 'account-balance'; + +export interface UsageCredential { + token: string; + method: 'api-key' | 'oauth'; +} + +export interface UsageEndpointSelection { + url: string; + kind: UsageEndpointKind; +} + +export function selectUsageEndpoint( + baseUrl: string, + credential: UsageCredential, +): UsageEndpointSelection { + if (credential.method === 'api-key' && isSecretApiKey(credential.token)) { + return { + url: accountBalanceEndpoint(baseUrl), + kind: 'account-balance', + }; + } + + return { + url: quotaEndpoint(baseUrl), + kind: 'quota', + }; +} + export function usageEndpoint(baseUrl: string, apiKey: string): string { - return isSecretApiKey(apiKey) - ? accountBalanceEndpoint(baseUrl) - : quotaEndpoint(baseUrl); + return selectUsageEndpoint(baseUrl, { token: apiKey, method: 'api-key' }).url; } export function fileUploadEndpoint(baseUrl: string): string { diff --git a/src/commands/quota/show.ts b/src/commands/quota/show.ts index e9210683..95bef531 100644 --- a/src/commands/quota/show.ts +++ b/src/commands/quota/show.ts @@ -1,15 +1,12 @@ import { defineCommand } from '../../command'; import { requestJson } from '../../client/http'; -import { quotaEndpoint, usageEndpoint } from '../../client/endpoints'; +import { selectUsageEndpoint } from '../../client/endpoints'; +import { resolveCredential } from '../../auth/resolver'; import { formatOutput, detectOutputFormat } from '../../output/formatter'; import { renderUsage } from '../../output/usage'; import type { Config } from '../../config/schema'; import type { GlobalFlags } from '../../types/flags'; -import type { AccountBalanceResponse, QuotaModelRemain } from '../../types/api'; - -interface QuotaApiResponse { - model_remains: QuotaModelRemain[]; -} +import type { AccountBalanceResponse, QuotaResponse } from '../../types/api'; export default defineCommand({ name: 'quota show', @@ -26,21 +23,20 @@ export default defineCommand({ } const format = detectOutputFormat(flags.output as string | undefined); - const apiKey = config.apiKey || config.fileApiKey; + const credential = await resolveCredential(config); + const endpoint = selectUsageEndpoint(config.baseUrl, credential); - if (apiKey && apiKey.startsWith('sk-api-')) { - const url = usageEndpoint(config.baseUrl, apiKey); - const response = await requestJson(config, { url }); + if (endpoint.kind === 'account-balance') { + const response = await requestJson(config, { url: endpoint.url }); if (format !== 'text') { console.log(formatOutput(response, format)); return; } - renderUsage(response, config, apiKey); + renderUsage(response, config, credential.token); return; } - const url = quotaEndpoint(config.baseUrl); - const response = await requestJson(config, { url }); + const response = await requestJson(config, { url: endpoint.url }); const models = response.model_remains || []; if (format !== 'text') { diff --git a/src/sdk/index.ts b/src/sdk/index.ts index fc46f3b5..9b9b6fb3 100644 --- a/src/sdk/index.ts +++ b/src/sdk/index.ts @@ -10,6 +10,16 @@ import { FileSDK } from "./file"; import { Client } from "./client"; import { MiniMaxSDKOptions } from "./types"; +export { + isAccountBalanceResponse, + isQuotaResponse, +} from "./quota"; +export type { + AccountBalanceResponse, + QuotaInfoResponse, + QuotaResponse, +} from "./quota"; + export class MiniMaxSDK extends Client { readonly text: TextSDK; readonly speech: SpeechSDK; diff --git a/src/sdk/quota/index.ts b/src/sdk/quota/index.ts index b48b359d..e4a23b40 100644 --- a/src/sdk/quota/index.ts +++ b/src/sdk/quota/index.ts @@ -1,12 +1,31 @@ import { Client } from "../client"; -import { quotaEndpoint } from "../../client/endpoints"; -import type { QuotaResponse } from "../../types/api"; +import { selectUsageEndpoint } from "../../client/endpoints"; +import { resolveCredential } from "../../auth/resolver"; +import type { AccountBalanceResponse, QuotaResponse } from "../../types/api"; + +export type QuotaInfoResponse = QuotaResponse | AccountBalanceResponse; + +export function isQuotaResponse(response: QuotaInfoResponse): response is QuotaResponse { + return 'model_remains' in response; +} + +export function isAccountBalanceResponse( + response: QuotaInfoResponse, +): response is AccountBalanceResponse { + return 'available_amount' in response; +} export class QuotaSDK extends Client { - async info(): Promise { - const url = quotaEndpoint(this.config.baseUrl); - const res = await this.requestJson({ url }); + async info(): Promise { + const credential = await resolveCredential(this.config); + const endpoint = selectUsageEndpoint(this.config.baseUrl, credential); + + if (endpoint.kind === 'account-balance') { + return this.requestJson({ url: endpoint.url }); + } - return res; + return this.requestJson({ url: endpoint.url }); } } + +export type { AccountBalanceResponse, QuotaResponse } from "../../types/api"; diff --git a/test/sdk/quota.test.ts b/test/sdk/quota.test.ts index f57cc407..57a67aa0 100644 --- a/test/sdk/quota.test.ts +++ b/test/sdk/quota.test.ts @@ -1,49 +1,203 @@ -import { describe, it, expect, mock } from 'bun:test'; -import { MiniMaxSDK } from '../../src/sdk'; +import { afterEach, describe, expect, it } from 'bun:test'; +import { mkdtempSync, rmSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { + isAccountBalanceResponse, + isQuotaResponse, + MiniMaxSDK, +} from '../../src/sdk'; +import { createMockServer, jsonResponse, type MockServer } from '../helpers/mock-server'; + +const quotaResponse = { + model_remains: [ + { + model_name: 'MiniMax-M3', + start_time: 0, + end_time: 9999999999, + remains_time: 1000, + current_interval_total_count: 1000, + current_interval_usage_count: 500, + current_interval_remaining_percent: 50, + current_weekly_total_count: 5000, + current_weekly_usage_count: 2000, + weekly_start_time: 0, + weekly_end_time: 9999999999, + weekly_remains_time: 3000, + }, + ], +}; + +const balanceResponse = { + available_amount: '98.00', + cash_balance: '0.00', + voucher_balance: '98.00', + credit_balance: '0.00', + owed_amount: '0.00', + balance_alert_switch: false, + balance_alert_threshold: '', + base_resp: { status_code: 0, status_msg: 'success' }, +}; describe('MiniMaxSDK.quota', () => { - it('should get quota info successfully', async () => { - const mockFetch = mock(async (url: string) => { - if (url.includes('/v1/token_plan/remains')) { - return new Response(JSON.stringify({ - model_remains: [ - { - model_name: 'MiniMax-M3', - start_time: 0, - end_time: 9999999999, - remains_time: 1000, - current_interval_total_count: 1000, - current_interval_usage_count: 500, - current_interval_remaining_percent: 50, - current_weekly_total_count: 5000, - current_weekly_usage_count: 2000, - weekly_start_time: 0, - weekly_end_time: 9999999999, - weekly_remains_time: 3000, - }, - ], - }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); - } - return new Response('Not found', { status: 404 }); + let server: MockServer | undefined; + let configDir: string | undefined; + const originalConfigDir = process.env.MMX_CONFIG_DIR; + + afterEach(() => { + server?.close(); + if (configDir) rmSync(configDir, { recursive: true, force: true }); + if (originalConfigDir === undefined) delete process.env.MMX_CONFIG_DIR; + else process.env.MMX_CONFIG_DIR = originalConfigDir; + server = undefined; + configDir = undefined; + }); + + function useIsolatedConfig(config: Record = {}): void { + configDir = mkdtempSync(join(tmpdir(), 'mmx-sdk-quota-')); + process.env.MMX_CONFIG_DIR = configDir; + writeFileSync(join(configDir, 'config.json'), JSON.stringify(config)); + } + + it('uses token_plan/remains for a regular API key from options', async () => { + const requests: Array<{ + url: string; + authorization: string | null; + xApiKey: string | null; + }> = []; + server = createMockServer({ + routes: { + '/v1/token_plan/remains': (req) => { + requests.push({ + url: req.url, + authorization: req.headers.get('Authorization'), + xApiKey: req.headers.get('x-api-key'), + }); + return jsonResponse(quotaResponse); + }, + }, + }); + useIsolatedConfig(); + + const sdk = new MiniMaxSDK({ apiKey: 'regular-api-key', baseUrl: server.url }); + const result = await sdk.quota.info(); + + expect(requests).toEqual([ + { + url: `${server.url}/v1/token_plan/remains`, + authorization: 'Bearer regular-api-key', + xApiKey: null, + }, + ]); + expect(isQuotaResponse(result)).toBe(true); + if (!isQuotaResponse(result)) throw new Error('Expected a quota response'); + expect(result.model_remains[0].model_name).toBe('MiniMax-M3'); + }); + + it('uses account/query_balance for an sk-api key from options', async () => { + const requests: Array<{ + url: string; + authorization: string | null; + xApiKey: string | null; + }> = []; + server = createMockServer({ + routes: { + '/account/query_balance': (req) => { + requests.push({ + url: req.url, + authorization: req.headers.get('Authorization'), + xApiKey: req.headers.get('x-api-key'), + }); + return jsonResponse(balanceResponse); + }, + }, + }); + useIsolatedConfig(); + + const sdk = new MiniMaxSDK({ apiKey: 'sk-api-secret-key', baseUrl: server.url }); + const result = await sdk.quota.info(); + + expect(requests).toEqual([ + { + url: `${server.url}/account/query_balance`, + authorization: 'Bearer sk-api-secret-key', + xApiKey: null, + }, + ]); + expect(isAccountBalanceResponse(result)).toBe(true); + if (!isAccountBalanceResponse(result)) throw new Error('Expected an account balance response'); + expect(result.available_amount).toBe('98.00'); + }); + + it('uses an sk-api key from the CLI config file', async () => { + const requests: Array<{ + url: string; + authorization: string | null; + xApiKey: string | null; + }> = []; + server = createMockServer({ + routes: { + '/account/query_balance': (req) => { + requests.push({ + url: req.url, + authorization: req.headers.get('Authorization'), + xApiKey: req.headers.get('x-api-key'), + }); + return jsonResponse(balanceResponse); + }, + }, }); + useIsolatedConfig({ api_key: 'sk-api-config-key' }); + + const sdk = new MiniMaxSDK({ baseUrl: server.url }); + const result = await sdk.quota.info(); - const originalFetch = globalThis.fetch; - globalThis.fetch = mockFetch as unknown as typeof fetch; + expect(requests).toEqual([ + { + url: `${server.url}/account/query_balance`, + authorization: 'Bearer sk-api-config-key', + xApiKey: null, + }, + ]); + expect(isAccountBalanceResponse(result)).toBe(true); + }); - try { - const sdk = new MiniMaxSDK({ - apiKey: 'test-key', - }); + it('uses token_plan/remains for OAuth even when its token has an sk-api prefix', async () => { + const requests: Array<{ + url: string; + authorization: string | null; + xApiKey: string | null; + }> = []; + server = createMockServer({ + routes: { + '/v1/token_plan/remains': (req) => { + requests.push({ + url: req.url, + authorization: req.headers.get('Authorization'), + xApiKey: req.headers.get('x-api-key'), + }); + return jsonResponse(quotaResponse); + }, + }, + }); + useIsolatedConfig({ + oauth: { + access_token: 'sk-api-oauth-token', + refresh_token: 'refresh-token', + expires_at: new Date(Date.now() + 60 * 60 * 1000).toISOString(), + }, + }); - const result = await sdk.quota.info(); + const sdk = new MiniMaxSDK({ baseUrl: server.url }); + const result = await sdk.quota.info(); - expect(result.model_remains).toHaveLength(1); - expect(result.model_remains[0].model_name).toBe('MiniMax-M3'); - } finally { - globalThis.fetch = originalFetch; - } + expect(requests).toEqual([ + { + url: `${server.url}/v1/token_plan/remains`, + authorization: 'Bearer sk-api-oauth-token', + xApiKey: null, + }, + ]); + expect(isQuotaResponse(result)).toBe(true); }); });