diff --git a/scripts/source-file-size-baseline.json b/scripts/source-file-size-baseline.json index 3415a93df1..634471d5a0 100644 --- a/scripts/source-file-size-baseline.json +++ b/scripts/source-file-size-baseline.json @@ -38,13 +38,13 @@ "src/tools/impl/enter-worktree.ts": 1020, "src/tools/manager.ts": 3016, "src/tools/tool-execution-context.test.ts": 1138, - "src/types/protocol_v2.ts": 2889, + "src/types/protocol_v2.ts": 2830, "src/websocket/listen-client-concurrency.test.ts": 2686, "src/websocket/listen-client-protocol.test.ts": 5820, "src/websocket/listener/commands/memory.ts": 1114, "src/websocket/listener/file-commands.ts": 1053, "src/websocket/listener/lifecycle.ts": 1051, - "src/websocket/listener/protocol-inbound.ts": 2264, + "src/websocket/listener/protocol-inbound.ts": 2256, "src/websocket/listener/protocol-outbound.ts": 1085, "src/websocket/listener/turn.ts": 1010 } diff --git a/src/providers/chatgpt-reset-credit-service.test.ts b/src/providers/chatgpt-reset-credit-service.test.ts new file mode 100644 index 0000000000..4187f3cc35 --- /dev/null +++ b/src/providers/chatgpt-reset-credit-service.test.ts @@ -0,0 +1,385 @@ +import { describe, expect, mock, test } from "bun:test"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { setLocalOAuthProvider } from "@/backend/local/local-provider-auth-store"; +import { + consumeChatGPTRateLimitResetCredit, + readChatGPTRateLimitResetCredits, +} from "@/providers/chatgpt-reset-credit-service"; + +async function withEnv( + updates: Record, + run: () => Promise, +): Promise { + const previous = Object.fromEntries( + Object.keys(updates).map((key) => [key, process.env[key]]), + ); + try { + for (const [key, value] of Object.entries(updates)) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + return await run(); + } finally { + for (const [key, value] of Object.entries(previous)) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + } +} + +function connectLocalProvider(input: { + storageDir: string; + providerName: string; +}): void { + setLocalOAuthProvider({ + storageDir: input.storageDir, + providerName: input.providerName, + providerType: "chatgpt_oauth", + auth: { + type: "oauth", + access: "chatgpt-access-token", + refresh: "chatgpt-refresh-token", + expires: Date.now() + 60_000, + accountId: "chatgpt-account-1", + }, + }); +} + +function resetCreditsResponse(availableCount: number): Response { + return Response.json({ + available_count: availableCount, + credits: + availableCount > 0 + ? [ + { + id: "RateLimitResetCredit_1", + reset_type: "codex_rate_limits", + status: "available", + granted_at: "2026-08-01T00:00:00Z", + expires_at: "2026-09-01T00:00:00Z", + title: "Full reset", + description: "Ready to redeem", + }, + ] + : [], + }); +} + +describe("ChatGPT reset-credit service", () => { + test("lists local reset credits with refreshed OAuth request credentials", async () => { + const storageDir = await mkdtemp(join(tmpdir(), "chatgpt-reset-local-")); + const providerName = "chatgpt-reset-local-list"; + try { + connectLocalProvider({ storageDir, providerName }); + const fetchSpy = mock(async (..._args: Parameters) => + resetCreditsResponse(1), + ); + const fetchMock = fetchSpy as unknown as typeof fetch; + + const result = await readChatGPTRateLimitResetCredits({ + target: "local", + providerName, + storageDir, + forceRefresh: true, + now: () => Date.parse("2026-08-06T12:00:00Z"), + fetch: fetchMock, + }); + + expect(result).toEqual({ + success: true, + credits: { + providerName, + fetchedAt: "2026-08-06T12:00:00.000Z", + availableCount: 1, + credits: [ + { + id: "RateLimitResetCredit_1", + resetType: "codex_rate_limits", + status: "available", + grantedAt: "2026-08-01T00:00:00Z", + expiresAt: "2026-09-01T00:00:00Z", + title: "Full reset", + description: "Ready to redeem", + }, + ], + }, + }); + const call = fetchSpy.mock.calls[0]; + expect(String(call?.[0])).toBe( + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits", + ); + const headers = new Headers(call?.[1]?.headers); + expect(headers.get("authorization")).toBe("Bearer chatgpt-access-token"); + expect(headers.get("chatgpt-account-id")).toBe("chatgpt-account-1"); + } finally { + await rm(storageDir, { recursive: true, force: true }); + } + }); + + test("lists API-stored reset credits through Letta Cloud and caches them", async () => { + const fetchSpy = mock(async (..._args: Parameters) => + Response.json({ + providerName: "chatgpt-reset-cloud-list", + fetchedAt: "2026-08-06T12:00:00.000Z", + availableCount: 1, + credits: [ + { + id: "RateLimitResetCredit_1", + resetType: "codex_rate_limits", + status: "available", + grantedAt: "2026-08-01T00:00:00Z", + expiresAt: null, + title: null, + description: null, + }, + ], + }), + ); + const fetchMock = fetchSpy as unknown as typeof fetch; + const input = { + target: "api" as const, + providerName: "chatgpt-reset-cloud-list", + fetch: fetchMock, + now: () => Date.parse("2026-08-06T12:00:00Z"), + getSettings: async () => ({ + env: { + LETTA_API_KEY: "letta-access-token", + LETTA_BASE_URL: "https://api.test.letta.com", + }, + refreshToken: undefined, + tokenExpiresAt: undefined, + }), + }; + + const [first, second] = await withEnv( + { LETTA_API_KEY: undefined, LETTA_BASE_URL: undefined }, + async () => [ + await readChatGPTRateLimitResetCredits(input), + await readChatGPTRateLimitResetCredits(input), + ], + ); + + expect(first).toEqual(second); + expect(fetchSpy).toHaveBeenCalledTimes(1); + const call = fetchSpy.mock.calls[0]; + expect(String(call?.[0])).toBe( + "https://api.test.letta.com/v1/providers/chatgpt-rate-limit-reset-credits?provider_name=chatgpt-reset-cloud-list", + ); + const headers = new Headers(call?.[1]?.headers); + expect(headers.get("authorization")).toBe("Bearer letta-access-token"); + }); + + test("passes the same idempotency key and selected reset ID to Cloud", async () => { + const fetchSpy = mock(async (...args: Parameters) => { + const [url] = args; + const value = String(url); + if (value.endsWith("/consume")) { + return Response.json({ outcome: "no_credit" }); + } + throw new Error(`Unexpected URL ${value}`); + }); + const fetchMock = fetchSpy as unknown as typeof fetch; + + const result = await withEnv( + { LETTA_API_KEY: undefined, LETTA_BASE_URL: undefined }, + () => + consumeChatGPTRateLimitResetCredit({ + target: "api", + providerName: "chatgpt-reset-cloud-consume", + idempotencyKey: "redeem-request-1", + resetId: "RateLimitResetCredit_1", + fetch: fetchMock, + getSettings: async () => ({ + env: { + LETTA_API_KEY: "letta-access-token", + LETTA_BASE_URL: "https://api.test.letta.com", + }, + refreshToken: undefined, + tokenExpiresAt: undefined, + }), + }), + ); + + expect(result).toEqual({ success: true, outcome: "no_credit" }); + const call = fetchSpy.mock.calls[0]; + expect(String(call?.[0])).toBe( + "https://api.test.letta.com/v1/providers/chatgpt-rate-limit-reset-credits/consume", + ); + expect(JSON.parse(String(call?.[1]?.body))).toEqual({ + provider_name: "chatgpt-reset-cloud-consume", + idempotency_key: "redeem-request-1", + reset_id: "RateLimitResetCredit_1", + }); + }); + + test("refreshes usage and reset inventory after a successful local reset", async () => { + const storageDir = await mkdtemp(join(tmpdir(), "chatgpt-reset-refresh-")); + const providerName = "chatgpt-reset-local-refresh"; + let listReads = 0; + try { + connectLocalProvider({ storageDir, providerName }); + const fetchMock = mock(async (url: Parameters[0]) => { + const value = String(url); + if (value.endsWith("/consume")) { + return Response.json({ code: "reset", windows_reset: 2 }); + } + if (value.endsWith("/wham/usage")) { + return Response.json({ + plan_type: "pro", + rate_limit: { + limit_reached: false, + primary_window: { + used_percent: 0, + limit_window_seconds: 18_000, + reset_after_seconds: 18_000, + }, + }, + }); + } + if (value.endsWith("/rate-limit-reset-credits")) { + listReads += 1; + return resetCreditsResponse(listReads === 1 ? 1 : 0); + } + throw new Error(`Unexpected URL ${value}`); + }) as unknown as typeof fetch; + const sharedInput = { + target: "local" as const, + providerName, + storageDir, + fetch: fetchMock, + now: () => Date.parse("2026-08-06T12:00:00Z"), + }; + + const before = await readChatGPTRateLimitResetCredits(sharedInput); + expect(before.success && before.credits.availableCount).toBe(1); + + const consumed = await consumeChatGPTRateLimitResetCredit({ + ...sharedInput, + idempotencyKey: "redeem-refresh-1", + }); + + expect(consumed.success).toBe(true); + if (!consumed.success) throw new Error(consumed.error.message); + expect(consumed.outcome).toBe("reset"); + expect(consumed.refreshedUsage?.limitReached).toBe(false); + expect(consumed.refreshedCredits?.availableCount).toBe(0); + expect(consumed.refreshError).toBeUndefined(); + expect(listReads).toBe(2); + + const after = await readChatGPTRateLimitResetCredits(sharedInput); + expect(after.success && after.credits.availableCount).toBe(0); + expect(listReads).toBe(2); + } finally { + await rm(storageDir, { recursive: true, force: true }); + } + }); + + test("treats already_redeemed as idempotent success and refreshes state", async () => { + const readUsage = mock(async () => ({ + success: true as const, + usage: { + providerName: "chatgpt-reset-idempotent", + fetchedAt: "2026-08-06T12:00:00.000Z", + summary: "Usage: no active quota window reported", + primary: null, + secondary: null, + additional: [], + }, + })); + let listCalls = 0; + const fetchMock = mock(async (url: Parameters[0]) => { + const value = String(url); + if (value.endsWith("/consume")) { + return Response.json({ outcome: "already_redeemed" }); + } + listCalls += 1; + return Response.json({ + providerName: "chatgpt-reset-idempotent", + fetchedAt: "2026-08-06T12:00:00.000Z", + availableCount: 0, + credits: [], + }); + }) as unknown as typeof fetch; + + const result = await withEnv( + { LETTA_API_KEY: undefined, LETTA_BASE_URL: undefined }, + () => + consumeChatGPTRateLimitResetCredit({ + target: "api", + providerName: "chatgpt-reset-idempotent", + idempotencyKey: "redeem-idempotent-1", + fetch: fetchMock, + readUsage, + getSettings: async () => ({ + env: { + LETTA_API_KEY: "letta-access-token", + LETTA_BASE_URL: "https://api.test.letta.com", + }, + refreshToken: undefined, + tokenExpiresAt: undefined, + }), + }), + ); + + expect(result.success && result.outcome).toBe("already_redeemed"); + expect(readUsage).toHaveBeenCalledTimes(1); + expect(listCalls).toBe(1); + }); + + test("does not refresh state for non-success consume outcomes", async () => { + const readUsage = mock(async () => { + throw new Error("usage refresh must not run"); + }); + const fetchMock = mock(async () => + Response.json({ outcome: "nothing_to_reset" }), + ) as unknown as typeof fetch; + + const result = await withEnv( + { LETTA_API_KEY: undefined, LETTA_BASE_URL: undefined }, + () => + consumeChatGPTRateLimitResetCredit({ + target: "api", + providerName: "chatgpt-reset-nothing", + idempotencyKey: "redeem-nothing-1", + fetch: fetchMock, + readUsage, + getSettings: async () => ({ + env: { + LETTA_API_KEY: "letta-access-token", + LETTA_BASE_URL: "https://api.test.letta.com", + }, + refreshToken: undefined, + tokenExpiresAt: undefined, + }), + }), + ); + + expect(result).toEqual({ success: true, outcome: "nothing_to_reset" }); + expect(readUsage).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + test("rejects an empty idempotency key before any request", async () => { + const fetchMock = mock(async () => + resetCreditsResponse(1), + ) as unknown as typeof fetch; + + const result = await consumeChatGPTRateLimitResetCredit({ + target: "api", + providerName: "chatgpt-reset-empty-key", + idempotencyKey: " ", + fetch: fetchMock, + }); + + expect(result).toEqual({ + success: false, + error: { + code: "bad_request", + message: "An idempotency key is required to consume a reset credit.", + }, + }); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/src/providers/chatgpt-reset-credit-service.ts b/src/providers/chatgpt-reset-credit-service.ts new file mode 100644 index 0000000000..777c729937 --- /dev/null +++ b/src/providers/chatgpt-reset-credit-service.ts @@ -0,0 +1,705 @@ +import { hostname } from "node:os"; +import { + LETTA_CLOUD_API_URL, + refreshAccessToken as refreshLettaAccessToken, + type TokenResponse, +} from "@/auth/oauth"; +import { getLettaCodeHeaders } from "@/backend/api/http-headers"; +import { + getLocalOAuthApiKey, + getLocalProviderRecordByName, + LOCAL_CHATGPT_PROVIDER_NAME, + type LocalProviderRecord, +} from "@/backend/local/local-provider-auth-store"; +import { + type ChatGPTUsageError, + type ChatGPTUsageSnapshot, + type ReadChatGPTUsageInput, + readChatGPTUsage, +} from "@/providers/chatgpt-usage-service"; +import { type Settings, settingsManager } from "@/settings-manager"; + +const CHATGPT_RESET_CREDITS_URL = + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits"; +const CHATGPT_RESET_CREDITS_CONSUME_URL = + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume"; +const CLOUD_RESET_CREDITS_PATH = + "/v1/providers/chatgpt-rate-limit-reset-credits"; +const CLOUD_RESET_CREDITS_CONSUME_PATH = + "/v1/providers/chatgpt-rate-limit-reset-credits/consume"; +const OPENAI_CODEX_OAUTH_PROVIDER_ID = "openai-codex"; +const CACHE_TTL_MS = 30_000; +const DEFAULT_TIMEOUT_MS = 15_000; +const TOKEN_REFRESH_BUFFER_MS = 5 * 60 * 1000; + +export interface ChatGPTRateLimitResetCredit { + id: string; + resetType: string; + status: string; + grantedAt: string; + expiresAt: string | null; + title: string | null; + description: string | null; +} + +export interface ChatGPTRateLimitResetCredits { + providerName: string; + fetchedAt: string; + availableCount: number; + credits: ChatGPTRateLimitResetCredit[]; +} + +export type ChatGPTRateLimitResetConsumeOutcome = + | "reset" + | "nothing_to_reset" + | "no_credit" + | "already_redeemed"; + +export type ChatGPTRateLimitResetCreditsReadResult = + | { success: true; credits: ChatGPTRateLimitResetCredits } + | { success: false; error: ChatGPTUsageError }; + +export type ChatGPTRateLimitResetCreditConsumeResult = + | { + success: true; + outcome: ChatGPTRateLimitResetConsumeOutcome; + refreshedUsage?: ChatGPTUsageSnapshot; + refreshedCredits?: ChatGPTRateLimitResetCredits; + refreshError?: ChatGPTUsageError; + } + | { success: false; error: ChatGPTUsageError }; + +export interface ReadChatGPTRateLimitResetCreditsInput + extends ReadChatGPTUsageInput {} + +export interface ConsumeChatGPTRateLimitResetCreditInput + extends Omit { + idempotencyKey: string; + resetId?: string; + readUsage?: typeof readChatGPTUsage; +} + +type JsonRecord = Record; + +type CachedResetCredits = { + expiresAt: number; + result: Extract; +}; + +type RequestContext = { + target: "local" | "api"; + providerName: string; + cacheKey: string; + listUrl: string; + consumeUrl: string; + headers: Record; +}; + +const resetCreditsCache = new Map(); + +function asRecord(value: unknown): JsonRecord | null { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as JsonRecord) + : null; +} + +function getValue( + record: JsonRecord | null | undefined, + keys: string[], +): unknown { + if (!record) return undefined; + for (const key of keys) { + const value = record[key]; + if (value !== undefined && value !== null) return value; + } + return undefined; +} + +function getString( + record: JsonRecord | null | undefined, + keys: string[], +): string | null { + const value = getValue(record, keys); + return typeof value === "string" && value.trim() ? value.trim() : null; +} + +function getNumber( + record: JsonRecord | null | undefined, + keys: string[], +): number | null { + const value = getValue(record, keys); + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : null; + } + return null; +} + +function resetCreditError( + code: ChatGPTUsageError["code"], + message: string, + retryAfter?: number, +): { success: false; error: ChatGPTUsageError } { + return { + success: false, + error: { + code, + message, + ...(retryAfter !== undefined ? { retryAfterMs: retryAfter } : {}), + }, + }; +} + +function retryAfterMs(response: Response): number | undefined { + const value = response.headers.get("retry-after"); + if (!value) return undefined; + const seconds = Number(value); + if (Number.isFinite(seconds)) return Math.max(0, seconds * 1000); + const retryAt = Date.parse(value); + return Number.isFinite(retryAt) + ? Math.max(0, retryAt - Date.now()) + : undefined; +} + +async function readJsonRecord(response: Response): Promise { + try { + return asRecord(await response.json()); + } catch { + return null; + } +} + +function responseMessage(raw: JsonRecord | null, fallback: string): string { + return getString(raw, ["message", "error", "detail"]) ?? fallback; +} + +function cloudBaseUrl(settings: Pick): string { + return ( + process.env.LETTA_BASE_URL || + settings.env?.LETTA_BASE_URL || + LETTA_CLOUD_API_URL + ).replace(/\/+$/, ""); +} + +async function cloudApiKey(input: { + settings: Pick; + now: number; + refreshAccessToken?: ( + refreshToken: string, + deviceId: string, + deviceName?: string, + ) => Promise; +}): Promise<{ apiKey: string | null; error?: ChatGPTUsageError }> { + const envApiKey = process.env.LETTA_API_KEY; + let apiKey = envApiKey || input.settings.env?.LETTA_API_KEY || null; + + if ( + !envApiKey && + input.settings.refreshToken && + (!apiKey || + (input.settings.tokenExpiresAt !== undefined && + input.settings.tokenExpiresAt - input.now < TOKEN_REFRESH_BUFFER_MS)) + ) { + try { + const refresh = input.refreshAccessToken ?? refreshLettaAccessToken; + const tokens = await refresh( + input.settings.refreshToken, + settingsManager.getOrCreateDeviceId(), + hostname(), + ); + apiKey = tokens.access_token; + settingsManager.updateSettings({ + env: { LETTA_API_KEY: tokens.access_token }, + refreshToken: tokens.refresh_token || input.settings.refreshToken, + tokenExpiresAt: input.now + tokens.expires_in * 1000, + }); + } catch { + return { + apiKey: null, + error: { + code: "refresh_failed", + message: "Failed to refresh the Letta Cloud access token.", + }, + }; + } + } + + return apiKey + ? { apiKey } + : { + apiKey: null, + error: { + code: "unauthorized", + message: "Sign in with Letta to read ChatGPT reset credits.", + }, + }; +} + +function localProviderNames(providerName: string | undefined): string[] { + return providerName?.trim() + ? [providerName.trim()] + : [LOCAL_CHATGPT_PROVIDER_NAME, OPENAI_CODEX_OAUTH_PROVIDER_ID]; +} + +function isConnectedChatGPTOAuthRecord( + record: LocalProviderRecord | null, +): record is LocalProviderRecord & { + auth: { type: "oauth"; accountId?: string }; +} { + return Boolean( + record && + record.auth.type === "oauth" && + (record.provider_type === "chatgpt_oauth" || + record.provider_type === OPENAI_CODEX_OAUTH_PROVIDER_ID), + ); +} + +async function resolveLocalContext( + input: ReadChatGPTRateLimitResetCreditsInput, +): Promise< + | { success: true; context: RequestContext } + | { success: false; error: ChatGPTUsageError } +> { + const record = localProviderNames(input.providerName) + .map((name) => getLocalProviderRecordByName(name, input.storageDir)) + .find(isConnectedChatGPTOAuthRecord); + if (!record) { + return resetCreditError( + "not_connected", + "No local ChatGPT OAuth provider is connected.", + ); + } + + let oauthApiKey: Awaited>; + try { + oauthApiKey = await getLocalOAuthApiKey({ + providerId: OPENAI_CODEX_OAUTH_PROVIDER_ID, + providerNames: [record.name], + storageDir: input.storageDir, + }); + } catch { + return resetCreditError( + "refresh_failed", + "Failed to refresh the ChatGPT OAuth token.", + ); + } + if (!oauthApiKey) { + return resetCreditError( + "not_connected", + "No local ChatGPT OAuth token is available.", + ); + } + + const accountId = + typeof oauthApiKey.credentials.accountId === "string" + ? oauthApiKey.credentials.accountId + : typeof record.auth.accountId === "string" + ? record.auth.accountId + : undefined; + return { + success: true, + context: { + target: "local", + providerName: record.name, + cacheKey: `local:${record.name}`, + listUrl: CHATGPT_RESET_CREDITS_URL, + consumeUrl: CHATGPT_RESET_CREDITS_CONSUME_URL, + headers: { + Accept: "application/json", + Authorization: `Bearer ${oauthApiKey.apiKey}`, + "User-Agent": "letta-code", + ...(accountId ? { "chatgpt-account-id": accountId } : {}), + }, + }, + }; +} + +async function resolveApiContext( + input: ReadChatGPTRateLimitResetCreditsInput, + now: number, +): Promise< + | { success: true; context: RequestContext } + | { success: false; error: ChatGPTUsageError } +> { + const providerName = input.providerName?.trim(); + if (!providerName) { + return resetCreditError( + "bad_request", + "A ChatGPT provider name is required for cloud reset credits.", + ); + } + + let settings: Pick; + try { + settings = await ( + input.getSettings ?? (() => settingsManager.getSettingsWithSecureTokens()) + )(); + } catch { + return resetCreditError( + "unauthorized", + "Failed to read Letta Cloud credentials.", + ); + } + + const baseUrl = cloudBaseUrl(settings); + const auth = await cloudApiKey({ + settings, + now, + refreshAccessToken: input.refreshAccessToken, + }); + if (auth.error || !auth.apiKey) { + return { + success: false, + error: auth.error ?? { + code: "unauthorized", + message: "Sign in with Letta to read ChatGPT reset credits.", + }, + }; + } + + const listUrl = new URL(`${baseUrl}${CLOUD_RESET_CREDITS_PATH}`); + listUrl.searchParams.set("provider_name", providerName); + return { + success: true, + context: { + target: "api", + providerName, + cacheKey: `api:${baseUrl}:${providerName}`, + listUrl: listUrl.toString(), + consumeUrl: `${baseUrl}${CLOUD_RESET_CREDITS_CONSUME_PATH}`, + headers: { + ...getLettaCodeHeaders(auth.apiKey), + Accept: "application/json", + }, + }, + }; +} + +async function resolveRequestContext( + input: ReadChatGPTRateLimitResetCreditsInput, + now: number, +): Promise< + | { success: true; context: RequestContext } + | { success: false; error: ChatGPTUsageError } +> { + const target = input.target ?? "local"; + if (target === "local") return resolveLocalContext(input); + if (target === "api") return resolveApiContext(input, now); + return resetCreditError( + "unsupported_target", + "ChatGPT reset credits are only available for local or cloud ChatGPT OAuth providers.", + ); +} + +async function requestJson(input: { + context: RequestContext; + fetchImpl?: typeof fetch; + timeoutMs?: number; + method: "GET" | "POST"; + url: string; + body?: JsonRecord; +}): Promise< + | { success: true; raw: JsonRecord } + | { success: false; error: ChatGPTUsageError } +> { + const controller = new AbortController(); + let didTimeOut = false; + const timeout = setTimeout(() => { + didTimeOut = true; + controller.abort(); + }, input.timeoutMs ?? DEFAULT_TIMEOUT_MS); + + try { + const response = await (input.fetchImpl ?? fetch)(input.url, { + method: input.method, + headers: { + ...input.context.headers, + ...(input.body ? { "Content-Type": "application/json" } : {}), + }, + ...(input.body ? { body: JSON.stringify(input.body) } : {}), + signal: controller.signal, + }); + + if (response.status === 401) { + return resetCreditError( + "unauthorized", + input.context.target === "api" + ? "Sign in with Letta to read ChatGPT reset credits." + : "ChatGPT rejected the OAuth token. Reconnect ChatGPT Plus/Pro and try again.", + ); + } + if (response.status === 403) { + return resetCreditError( + "forbidden", + "ChatGPT reset credits are not available for this account.", + ); + } + if (response.status === 400 && input.context.target === "api") { + const raw = await readJsonRecord(response); + return resetCreditError( + "bad_request", + responseMessage(raw, "The reset-credit request was invalid."), + ); + } + if (response.status === 404 && input.context.target === "api") { + const raw = await readJsonRecord(response); + return raw + ? resetCreditError( + "not_connected", + responseMessage( + raw, + "No cloud ChatGPT OAuth provider is connected.", + ), + ) + : resetCreditError( + "network_error", + "Letta Cloud ChatGPT reset-credit endpoint is unavailable.", + ); + } + if (response.status === 429) { + const raw = + input.context.target === "api" ? await readJsonRecord(response) : null; + return resetCreditError( + "rate_limited", + raw + ? responseMessage( + raw, + "ChatGPT reset credits are rate limited. Try again later.", + ) + : "ChatGPT reset credits are rate limited. Try again later.", + getNumber(raw, ["retryAfterMs", "retry_after_ms"]) ?? + retryAfterMs(response), + ); + } + if (!response.ok) { + const raw = + input.context.target === "api" ? await readJsonRecord(response) : null; + return resetCreditError( + "network_error", + raw + ? responseMessage( + raw, + `Letta Cloud ChatGPT reset-credit request failed with HTTP ${response.status}.`, + ) + : `ChatGPT reset-credit request failed with HTTP ${response.status}.`, + ); + } + + const raw = await readJsonRecord(response); + return raw + ? { success: true, raw } + : resetCreditError( + didTimeOut ? "network_error" : "bad_response", + didTimeOut + ? "ChatGPT reset-credit request timed out." + : "ChatGPT reset-credit request returned invalid JSON.", + ); + } catch { + return resetCreditError( + "network_error", + didTimeOut + ? "ChatGPT reset-credit request timed out." + : "Failed to fetch ChatGPT reset credits.", + ); + } finally { + clearTimeout(timeout); + } +} + +function normalizeCredit(value: unknown): ChatGPTRateLimitResetCredit | null { + const raw = asRecord(value); + const id = getString(raw, ["id"]); + const resetType = getString(raw, ["reset_type", "resetType"]); + const status = getString(raw, ["status"]); + const grantedAt = getString(raw, ["granted_at", "grantedAt"]); + if (!id || !resetType || !status || !grantedAt) return null; + return { + id, + resetType, + status, + grantedAt, + expiresAt: getString(raw, ["expires_at", "expiresAt"]), + title: getString(raw, ["title"]), + description: getString(raw, ["description"]), + }; +} + +function normalizeCredits(input: { + raw: JsonRecord; + providerName: string; + now: number; +}): ChatGPTRateLimitResetCredits | null { + const availableCount = getNumber(input.raw, [ + "available_count", + "availableCount", + ]); + const rawCredits = input.raw.credits; + if ( + availableCount === null || + !Number.isInteger(availableCount) || + availableCount < 0 || + !Array.isArray(rawCredits) + ) { + return null; + } + const credits = rawCredits.map(normalizeCredit); + if (credits.some((credit) => credit === null)) return null; + return { + providerName: input.providerName, + fetchedAt: new Date(input.now).toISOString(), + availableCount, + credits: credits as ChatGPTRateLimitResetCredit[], + }; +} + +function normalizeOutcome( + raw: JsonRecord, +): ChatGPTRateLimitResetConsumeOutcome | null { + const outcome = getString(raw, ["outcome", "code"]); + if (outcome === "nothingToReset") return "nothing_to_reset"; + if (outcome === "noCredit") return "no_credit"; + if (outcome === "alreadyRedeemed") return "already_redeemed"; + if ( + outcome === "reset" || + outcome === "nothing_to_reset" || + outcome === "no_credit" || + outcome === "already_redeemed" + ) { + return outcome; + } + return null; +} + +export async function readChatGPTRateLimitResetCredits( + input: ReadChatGPTRateLimitResetCreditsInput = {}, +): Promise { + const now = input.now?.() ?? Date.now(); + const resolved = await resolveRequestContext(input, now); + if (!resolved.success) return resolved; + + const cached = resetCreditsCache.get(resolved.context.cacheKey); + if (!input.forceRefresh && cached && cached.expiresAt > now) { + return cached.result; + } + + const response = await requestJson({ + context: resolved.context, + fetchImpl: input.fetch, + timeoutMs: input.timeoutMs, + method: "GET", + url: resolved.context.listUrl, + }); + if (!response.success) return response; + + const credits = normalizeCredits({ + raw: response.raw, + providerName: resolved.context.providerName, + now, + }); + if (!credits) { + return resetCreditError( + "bad_response", + "ChatGPT reset credits returned an invalid payload.", + ); + } + const result = { success: true as const, credits }; + resetCreditsCache.set(resolved.context.cacheKey, { + expiresAt: now + CACHE_TTL_MS, + result, + }); + return result; +} + +async function refreshRateLimitState( + input: ConsumeChatGPTRateLimitResetCreditInput, +): Promise<{ + refreshedUsage?: ChatGPTUsageSnapshot; + refreshedCredits?: ChatGPTRateLimitResetCredits; + refreshError?: ChatGPTUsageError; +}> { + const refreshInput: ReadChatGPTUsageInput = { + target: input.target, + providerName: input.providerName, + forceRefresh: true, + storageDir: input.storageDir, + timeoutMs: input.timeoutMs, + fetch: input.fetch, + now: input.now, + getSettings: input.getSettings, + refreshAccessToken: input.refreshAccessToken, + }; + const [usage, credits] = await Promise.all([ + (input.readUsage ?? readChatGPTUsage)(refreshInput), + readChatGPTRateLimitResetCredits(refreshInput), + ]); + const refreshError = !usage.success + ? usage.error + : !credits.success + ? credits.error + : undefined; + return { + ...(usage.success ? { refreshedUsage: usage.usage } : {}), + ...(credits.success ? { refreshedCredits: credits.credits } : {}), + ...(refreshError ? { refreshError } : {}), + }; +} + +export async function consumeChatGPTRateLimitResetCredit( + input: ConsumeChatGPTRateLimitResetCreditInput, +): Promise { + const idempotencyKey = input.idempotencyKey.trim(); + const resetId = input.resetId?.trim(); + if (!idempotencyKey) { + return resetCreditError( + "bad_request", + "An idempotency key is required to consume a reset credit.", + ); + } + if (input.resetId !== undefined && !resetId) { + return resetCreditError("bad_request", "Reset ID cannot be empty."); + } + + const now = input.now?.() ?? Date.now(); + const resolved = await resolveRequestContext(input, now); + if (!resolved.success) return resolved; + + const body = + resolved.context.target === "local" + ? { + redeem_request_id: idempotencyKey, + ...(resetId ? { credit_id: resetId } : {}), + } + : { + provider_name: resolved.context.providerName, + idempotency_key: idempotencyKey, + ...(resetId ? { reset_id: resetId } : {}), + }; + const response = await requestJson({ + context: resolved.context, + fetchImpl: input.fetch, + timeoutMs: input.timeoutMs, + method: "POST", + url: resolved.context.consumeUrl, + body, + }); + if (!response.success) return response; + + const outcome = normalizeOutcome(response.raw); + if (!outcome) { + return resetCreditError( + "bad_response", + "ChatGPT reset-credit consumption returned an invalid payload.", + ); + } + if (outcome !== "reset" && outcome !== "already_redeemed") { + return { success: true, outcome }; + } + + resetCreditsCache.delete(resolved.context.cacheKey); + return { + success: true, + outcome, + ...(await refreshRateLimitState(input)), + }; +} diff --git a/src/types/chatgpt-usage-protocol.ts b/src/types/chatgpt-usage-protocol.ts new file mode 100644 index 0000000000..e2af824cf0 --- /dev/null +++ b/src/types/chatgpt-usage-protocol.ts @@ -0,0 +1,129 @@ +export type ChatGPTUsageReadTarget = "local" | "api"; + +export interface ChatGPTUsageReadCommand { + type: "chatgpt_usage_read"; + request_id: string; + target: ChatGPTUsageReadTarget; + provider_name?: string; + force_refresh?: boolean; +} + +export interface ChatGPTRateLimitResetCreditsListCommand { + type: "chatgpt_rate_limit_reset_credits_list"; + request_id: string; + target: ChatGPTUsageReadTarget; + provider_name?: string; + force_refresh?: boolean; +} + +export interface ChatGPTRateLimitResetCreditConsumeCommand { + type: "chatgpt_rate_limit_reset_credit_consume"; + request_id: string; + target: ChatGPTUsageReadTarget; + provider_name?: string; + idempotency_key: string; + reset_id?: string; +} + +export interface ChatGPTUsageWindowPayload { + label: string; + usedPercent: number | null; + windowDurationMins: number | null; + resetsAt: number | null; +} + +export interface ChatGPTUsageCreditsPayload { + balance?: string | null; + availableCount?: number | null; + hasCredits?: boolean | null; + unlimited?: boolean | null; +} + +export interface ChatGPTUsageIndividualLimitPayload { + limit: string; + used: string; + remainingPercent: number; + resetsAt: number; +} + +export interface ChatGPTUsageSnapshotPayload { + providerName: string; + fetchedAt: string; + summary: string; + planType?: string | null; + limitReached?: boolean | null; + rateLimitReachedType?: string | null; + primary: ChatGPTUsageWindowPayload | null; + secondary: ChatGPTUsageWindowPayload | null; + additional: ChatGPTUsageWindowPayload[]; + credits?: ChatGPTUsageCreditsPayload | null; + individualLimit?: ChatGPTUsageIndividualLimitPayload | null; +} + +export interface ChatGPTRateLimitResetCreditPayload { + id: string; + resetType: string; + status: string; + grantedAt: string; + expiresAt: string | null; + title: string | null; + description: string | null; +} + +export interface ChatGPTRateLimitResetCreditsPayload { + providerName: string; + fetchedAt: string; + availableCount: number; + credits: ChatGPTRateLimitResetCreditPayload[]; +} + +export interface ChatGPTUsageReadErrorPayload { + code: + | "bad_request" + | "not_connected" + | "unsupported_target" + | "refresh_failed" + | "unauthorized" + | "forbidden" + | "rate_limited" + | "network_error" + | "bad_response"; + message: string; + retryAfterMs?: number; +} + +export interface ChatGPTUsageReadResponseMessage { + type: "chatgpt_usage_read_response"; + request_id: string; + success: boolean; + target: ChatGPTUsageReadTarget; + usage?: ChatGPTUsageSnapshotPayload; + error?: ChatGPTUsageReadErrorPayload; +} + +export interface ChatGPTRateLimitResetCreditsListResponseMessage { + type: "chatgpt_rate_limit_reset_credits_list_response"; + request_id: string; + success: boolean; + target: ChatGPTUsageReadTarget; + credits?: ChatGPTRateLimitResetCreditsPayload; + error?: ChatGPTUsageReadErrorPayload; +} + +export type ChatGPTRateLimitResetConsumeOutcome = + | "reset" + | "nothing_to_reset" + | "no_credit" + | "already_redeemed"; + +export interface ChatGPTRateLimitResetCreditConsumeResponseMessage { + type: "chatgpt_rate_limit_reset_credit_consume_response"; + request_id: string; + success: boolean; + target: ChatGPTUsageReadTarget; + outcome?: ChatGPTRateLimitResetConsumeOutcome; + refreshed_usage?: ChatGPTUsageSnapshotPayload; + refreshed_credits?: ChatGPTRateLimitResetCreditsPayload; + refresh_error?: ChatGPTUsageReadErrorPayload; + error?: ChatGPTUsageReadErrorPayload; +} diff --git a/src/types/protocol_v2.ts b/src/types/protocol_v2.ts index ce721f63f7..ed8a4687bb 100644 --- a/src/types/protocol_v2.ts +++ b/src/types/protocol_v2.ts @@ -35,6 +35,14 @@ import type { AppServerInfoResponseMessage, } from "./app-server-info"; import type { BackgroundProcessSummary } from "./background-process-protocol"; +import type { + ChatGPTRateLimitResetCreditConsumeCommand, + ChatGPTRateLimitResetCreditConsumeResponseMessage, + ChatGPTRateLimitResetCreditsListCommand, + ChatGPTRateLimitResetCreditsListResponseMessage, + ChatGPTUsageReadCommand, + ChatGPTUsageReadResponseMessage, +} from "./chatgpt-usage-protocol"; import type { ConversationForkBody } from "./conversation-fork-protocol"; import type { ExternalToolCallRequestMessage, @@ -47,6 +55,7 @@ import type { RuntimeScope } from "./runtime-scope"; import type { CronRunLogPage, CronTask } from "./schedule-protocol"; export type * from "./background-process-protocol"; +export type * from "./chatgpt-usage-protocol"; export type * from "./external-tool-protocol"; export type * from "./runtime-scope"; export type * from "./schedule-protocol"; @@ -1333,7 +1342,6 @@ export interface ListModelsCommand { } export type ConnectProviderStorageTarget = "local"; -export type ChatGPTUsageReadTarget = "local" | "api"; export interface ListConnectProvidersCommand { type: "list_connect_providers"; @@ -1369,18 +1377,6 @@ export interface DisconnectProviderCommand { provider_name?: string; } -export interface ChatGPTUsageReadCommand { - type: "chatgpt_usage_read"; - /** Echoed back in the response for request correlation. */ - request_id: string; - /** Provider store to inspect. */ - target: ChatGPTUsageReadTarget; - /** Optional connected ChatGPT provider alias. Defaults to the built-in alias. */ - provider_name?: string; - /** Skip the short listener-side cache. */ - force_refresh?: boolean; -} - export interface ConnectProviderField { key: string; label: string; @@ -1454,65 +1450,6 @@ export interface DisconnectProviderResponseMessage { error?: string; } -export interface ChatGPTUsageWindowPayload { - label: string; - usedPercent: number | null; - windowDurationMins: number | null; - resetsAt: number | null; -} - -export interface ChatGPTUsageCreditsPayload { - balance?: string | null; - availableCount?: number | null; - hasCredits?: boolean | null; - unlimited?: boolean | null; -} - -export interface ChatGPTUsageIndividualLimitPayload { - limit: string; - used: string; - remainingPercent: number; - resetsAt: number; -} - -export interface ChatGPTUsageSnapshotPayload { - providerName: string; - fetchedAt: string; - summary: string; - planType?: string | null; - limitReached?: boolean | null; - rateLimitReachedType?: string | null; - primary: ChatGPTUsageWindowPayload | null; - secondary: ChatGPTUsageWindowPayload | null; - additional: ChatGPTUsageWindowPayload[]; - credits?: ChatGPTUsageCreditsPayload | null; - individualLimit?: ChatGPTUsageIndividualLimitPayload | null; -} - -export interface ChatGPTUsageReadErrorPayload { - code: - | "bad_request" - | "not_connected" - | "unsupported_target" - | "refresh_failed" - | "unauthorized" - | "forbidden" - | "rate_limited" - | "network_error" - | "bad_response"; - message: string; - retryAfterMs?: number; -} - -export interface ChatGPTUsageReadResponseMessage { - type: "chatgpt_usage_read_response"; - request_id: string; - success: boolean; - target: ChatGPTUsageReadTarget; - usage?: ChatGPTUsageSnapshotPayload; - error?: ChatGPTUsageReadErrorPayload; -} - export interface UpdateModelPayload { /** Preferred model identifier from models.json (e.g. "sonnet") */ model_id?: string; @@ -2717,6 +2654,8 @@ export type WsProtocolCommand = | ConnectProviderCommand | DisconnectProviderCommand | ChatGPTUsageReadCommand + | ChatGPTRateLimitResetCreditsListCommand + | ChatGPTRateLimitResetCreditConsumeCommand | UpdateModelCommand | UpdateToolsetCommand | CronListCommand @@ -2818,6 +2757,8 @@ export type WsProtocolMessage = | ConnectProviderResponseMessage | DisconnectProviderResponseMessage | ChatGPTUsageReadResponseMessage + | ChatGPTRateLimitResetCreditsListResponseMessage + | ChatGPTRateLimitResetCreditConsumeResponseMessage | UpdateModelResponseMessage | UpdateToolsetResponseMessage | CronListResponseMessage diff --git a/src/websocket/listener/chatgpt-usage-protocol-inbound.test.ts b/src/websocket/listener/chatgpt-usage-protocol-inbound.test.ts new file mode 100644 index 0000000000..52877ad74b --- /dev/null +++ b/src/websocket/listener/chatgpt-usage-protocol-inbound.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, test } from "bun:test"; +import { + isChatGPTRateLimitResetCreditConsumeCommand, + isChatGPTRateLimitResetCreditsListCommand, + isChatGPTUsageReadCommand, +} from "@/websocket/listener/chatgpt-usage-protocol-inbound"; +import { parseServerMessage } from "@/websocket/listener/protocol-inbound"; + +describe("ChatGPT usage protocol input", () => { + test("preserves the existing usage read command", () => { + expect( + isChatGPTUsageReadCommand({ + type: "chatgpt_usage_read", + request_id: "usage-1", + target: "local", + force_refresh: true, + }), + ).toBe(true); + }); + + test("parses reset-credit list commands", () => { + const command = { + type: "chatgpt_rate_limit_reset_credits_list", + request_id: "credits-1", + target: "api", + provider_name: "chatgpt-pro", + force_refresh: true, + } as const; + + expect(isChatGPTRateLimitResetCreditsListCommand(command)).toBe(true); + expect(parseServerMessage(Buffer.from(JSON.stringify(command)))).toEqual( + command, + ); + }); + + test("parses reset-credit consume commands", () => { + const command = { + type: "chatgpt_rate_limit_reset_credit_consume", + request_id: "consume-1", + target: "local", + idempotency_key: "redeem-1", + reset_id: "credit-1", + } as const; + + expect(isChatGPTRateLimitResetCreditConsumeCommand(command)).toBe(true); + expect(parseServerMessage(Buffer.from(JSON.stringify(command)))).toEqual( + command, + ); + }); + + test("rejects invalid reset-credit commands", () => { + expect( + isChatGPTRateLimitResetCreditConsumeCommand({ + type: "chatgpt_rate_limit_reset_credit_consume", + request_id: "consume-1", + target: "local", + idempotency_key: " ", + }), + ).toBe(false); + expect( + isChatGPTRateLimitResetCreditsListCommand({ + type: "chatgpt_rate_limit_reset_credits_list", + request_id: "credits-1", + target: "workspace", + }), + ).toBe(false); + }); +}); diff --git a/src/websocket/listener/chatgpt-usage-protocol-inbound.ts b/src/websocket/listener/chatgpt-usage-protocol-inbound.ts new file mode 100644 index 0000000000..f1e523f192 --- /dev/null +++ b/src/websocket/listener/chatgpt-usage-protocol-inbound.ts @@ -0,0 +1,82 @@ +import type { + ChatGPTRateLimitResetCreditConsumeCommand, + ChatGPTRateLimitResetCreditsListCommand, + ChatGPTUsageReadCommand, +} from "@/types/chatgpt-usage-protocol"; + +function hasOptionalProviderName(value: { provider_name?: unknown }): boolean { + return ( + value.provider_name === undefined || typeof value.provider_name === "string" + ); +} + +function hasUsageTarget(value: { target?: unknown }): boolean { + return value.target === "local" || value.target === "api"; +} + +export function isChatGPTUsageReadCommand( + value: unknown, +): value is ChatGPTUsageReadCommand { + if (!value || typeof value !== "object") return false; + const command = value as { + type?: unknown; + request_id?: unknown; + target?: unknown; + provider_name?: unknown; + force_refresh?: unknown; + }; + return ( + command.type === "chatgpt_usage_read" && + typeof command.request_id === "string" && + hasUsageTarget(command) && + hasOptionalProviderName(command) && + (command.force_refresh === undefined || + typeof command.force_refresh === "boolean") + ); +} + +export function isChatGPTRateLimitResetCreditsListCommand( + value: unknown, +): value is ChatGPTRateLimitResetCreditsListCommand { + if (!value || typeof value !== "object") return false; + const command = value as { + type?: unknown; + request_id?: unknown; + target?: unknown; + provider_name?: unknown; + force_refresh?: unknown; + }; + return ( + command.type === "chatgpt_rate_limit_reset_credits_list" && + typeof command.request_id === "string" && + hasUsageTarget(command) && + hasOptionalProviderName(command) && + (command.force_refresh === undefined || + typeof command.force_refresh === "boolean") + ); +} + +export function isChatGPTRateLimitResetCreditConsumeCommand( + value: unknown, +): value is ChatGPTRateLimitResetCreditConsumeCommand { + if (!value || typeof value !== "object") return false; + const command = value as { + type?: unknown; + request_id?: unknown; + target?: unknown; + provider_name?: unknown; + idempotency_key?: unknown; + reset_id?: unknown; + }; + return ( + command.type === "chatgpt_rate_limit_reset_credit_consume" && + typeof command.request_id === "string" && + hasUsageTarget(command) && + hasOptionalProviderName(command) && + typeof command.idempotency_key === "string" && + command.idempotency_key.trim().length > 0 && + (command.reset_id === undefined || + (typeof command.reset_id === "string" && + command.reset_id.trim().length > 0)) + ); +} diff --git a/src/websocket/listener/commands/chatgpt-reset-credits.test.ts b/src/websocket/listener/commands/chatgpt-reset-credits.test.ts new file mode 100644 index 0000000000..6990e41c10 --- /dev/null +++ b/src/websocket/listener/commands/chatgpt-reset-credits.test.ts @@ -0,0 +1,141 @@ +import { describe, expect, test } from "bun:test"; +import type { + ConsumeChatGPTRateLimitResetCreditInput, + ReadChatGPTRateLimitResetCreditsInput, +} from "@/providers/chatgpt-reset-credit-service"; +import { + buildChatGPTResetCreditConsumeResponse, + buildChatGPTResetCreditsListResponse, +} from "@/websocket/listener/commands/chatgpt-reset-credits"; + +const credits = { + providerName: "chatgpt-pro", + fetchedAt: "2026-08-06T12:00:00.000Z", + availableCount: 1, + credits: [ + { + id: "credit-1", + resetType: "weekly", + status: "available", + grantedAt: "2026-08-01T12:00:00.000Z", + expiresAt: null, + title: "Weekly reset", + description: null, + }, + ], +}; + +describe("ChatGPT reset-credit listener commands", () => { + test("builds list responses and forwards the refresh flag", async () => { + let received: ReadChatGPTRateLimitResetCreditsInput | undefined; + const response = await buildChatGPTResetCreditsListResponse( + { + type: "chatgpt_rate_limit_reset_credits_list", + request_id: "list-1", + target: "api", + provider_name: "chatgpt-pro", + force_refresh: true, + }, + { + readCredits: async (input) => { + received = input; + return { success: true, credits }; + }, + }, + ); + + expect(received).toMatchObject({ + target: "api", + providerName: "chatgpt-pro", + forceRefresh: true, + }); + expect(response).toEqual({ + type: "chatgpt_rate_limit_reset_credits_list_response", + request_id: "list-1", + success: true, + target: "api", + credits, + }); + }); + + test("forwards idempotency fields and includes refreshed state", async () => { + let received: ConsumeChatGPTRateLimitResetCreditInput | undefined; + const refreshedUsage = { + providerName: "chatgpt-pro", + fetchedAt: "2026-08-06T12:00:01.000Z", + summary: "100% left", + primary: null, + secondary: null, + additional: [], + }; + const response = await buildChatGPTResetCreditConsumeResponse( + { + type: "chatgpt_rate_limit_reset_credit_consume", + request_id: "consume-1", + target: "local", + provider_name: "chatgpt-pro", + idempotency_key: "redeem-1", + reset_id: "credit-1", + }, + { + consumeCredit: async (input) => { + received = input; + return { + success: true, + outcome: "reset", + refreshedUsage, + refreshedCredits: credits, + }; + }, + }, + ); + + expect(received).toMatchObject({ + target: "local", + providerName: "chatgpt-pro", + idempotencyKey: "redeem-1", + resetId: "credit-1", + }); + expect(response).toMatchObject({ + type: "chatgpt_rate_limit_reset_credit_consume_response", + request_id: "consume-1", + success: true, + outcome: "reset", + refreshed_usage: refreshedUsage, + refreshed_credits: credits, + }); + }); + + test("returns structured service errors", async () => { + const response = await buildChatGPTResetCreditConsumeResponse( + { + type: "chatgpt_rate_limit_reset_credit_consume", + request_id: "consume-2", + target: "api", + idempotency_key: "redeem-2", + }, + { + consumeCredit: async () => ({ + success: false, + error: { + code: "rate_limited", + message: "Try again later.", + retryAfterMs: 1_000, + }, + }), + }, + ); + + expect(response).toEqual({ + type: "chatgpt_rate_limit_reset_credit_consume_response", + request_id: "consume-2", + success: false, + target: "api", + error: { + code: "rate_limited", + message: "Try again later.", + retryAfterMs: 1_000, + }, + }); + }); +}); diff --git a/src/websocket/listener/commands/chatgpt-reset-credits.ts b/src/websocket/listener/commands/chatgpt-reset-credits.ts new file mode 100644 index 0000000000..7b671b71df --- /dev/null +++ b/src/websocket/listener/commands/chatgpt-reset-credits.ts @@ -0,0 +1,165 @@ +import type WebSocket from "ws"; +import { + consumeChatGPTRateLimitResetCredit, + readChatGPTRateLimitResetCredits, +} from "@/providers/chatgpt-reset-credit-service"; +import type { + ChatGPTRateLimitResetCreditConsumeCommand, + ChatGPTRateLimitResetCreditConsumeResponseMessage, + ChatGPTRateLimitResetCreditsListCommand, + ChatGPTRateLimitResetCreditsListResponseMessage, +} from "@/types/chatgpt-usage-protocol"; +import { + isChatGPTRateLimitResetCreditConsumeCommand, + isChatGPTRateLimitResetCreditsListCommand, +} from "@/websocket/listener/chatgpt-usage-protocol-inbound"; +import type { RunDetachedListenerTask, SafeSocketSend } from "./types"; + +type ChatGPTResetCreditsCommandContext = { + socket: WebSocket; + safeSocketSend: SafeSocketSend; + runDetachedListenerTask: RunDetachedListenerTask; +}; + +type ChatGPTResetCreditsDependencies = { + readCredits?: typeof readChatGPTRateLimitResetCredits; + consumeCredit?: typeof consumeChatGPTRateLimitResetCredit; +}; + +export async function buildChatGPTResetCreditsListResponse( + command: ChatGPTRateLimitResetCreditsListCommand, + dependencies: ChatGPTResetCreditsDependencies = {}, +): Promise { + const result = await ( + dependencies.readCredits ?? readChatGPTRateLimitResetCredits + )({ + target: command.target, + ...(command.provider_name ? { providerName: command.provider_name } : {}), + forceRefresh: command.force_refresh === true, + }); + + if (!result.success) { + return { + type: "chatgpt_rate_limit_reset_credits_list_response", + request_id: command.request_id, + success: false, + target: command.target, + error: result.error, + }; + } + + return { + type: "chatgpt_rate_limit_reset_credits_list_response", + request_id: command.request_id, + success: true, + target: command.target, + credits: result.credits, + }; +} + +export async function buildChatGPTResetCreditConsumeResponse( + command: ChatGPTRateLimitResetCreditConsumeCommand, + dependencies: ChatGPTResetCreditsDependencies = {}, +): Promise { + const result = await ( + dependencies.consumeCredit ?? consumeChatGPTRateLimitResetCredit + )({ + target: command.target, + ...(command.provider_name ? { providerName: command.provider_name } : {}), + idempotencyKey: command.idempotency_key, + ...(command.reset_id ? { resetId: command.reset_id } : {}), + }); + + if (!result.success) { + return { + type: "chatgpt_rate_limit_reset_credit_consume_response", + request_id: command.request_id, + success: false, + target: command.target, + error: result.error, + }; + } + + return { + type: "chatgpt_rate_limit_reset_credit_consume_response", + request_id: command.request_id, + success: true, + target: command.target, + outcome: result.outcome, + ...(result.refreshedUsage + ? { refreshed_usage: result.refreshedUsage } + : {}), + ...(result.refreshedCredits + ? { refreshed_credits: result.refreshedCredits } + : {}), + ...(result.refreshError ? { refresh_error: result.refreshError } : {}), + }; +} + +function sendUnexpectedFailure( + socket: WebSocket, + command: + | ChatGPTRateLimitResetCreditsListCommand + | ChatGPTRateLimitResetCreditConsumeCommand, + safeSocketSend: SafeSocketSend, +): void { + const isList = command.type === "chatgpt_rate_limit_reset_credits_list"; + safeSocketSend( + socket, + { + type: isList + ? "chatgpt_rate_limit_reset_credits_list_response" + : "chatgpt_rate_limit_reset_credit_consume_response", + request_id: command.request_id, + success: false, + target: command.target, + error: { + code: "network_error", + message: "Failed to process the ChatGPT reset-credit request.", + }, + }, + "listener_chatgpt_reset_credits_send_failed", + "listener_chatgpt_reset_credits", + ); +} + +export function handleChatGPTResetCreditsCommand( + parsed: unknown, + context: ChatGPTResetCreditsCommandContext, +): boolean { + if (isChatGPTRateLimitResetCreditsListCommand(parsed)) { + const { socket, safeSocketSend, runDetachedListenerTask } = context; + runDetachedListenerTask("chatgpt_reset_credits_list", async () => { + try { + safeSocketSend( + socket, + await buildChatGPTResetCreditsListResponse(parsed), + "listener_chatgpt_reset_credits_send_failed", + "listener_chatgpt_reset_credits_list", + ); + } catch { + sendUnexpectedFailure(socket, parsed, safeSocketSend); + } + }); + return true; + } + + if (isChatGPTRateLimitResetCreditConsumeCommand(parsed)) { + const { socket, safeSocketSend, runDetachedListenerTask } = context; + runDetachedListenerTask("chatgpt_reset_credit_consume", async () => { + try { + safeSocketSend( + socket, + await buildChatGPTResetCreditConsumeResponse(parsed), + "listener_chatgpt_reset_credits_send_failed", + "listener_chatgpt_reset_credit_consume", + ); + } catch { + sendUnexpectedFailure(socket, parsed, safeSocketSend); + } + }); + return true; + } + + return false; +} diff --git a/src/websocket/listener/commands/chatgpt-usage.ts b/src/websocket/listener/commands/chatgpt-usage.ts index 3529446d04..a6dd8b2bbb 100644 --- a/src/websocket/listener/commands/chatgpt-usage.ts +++ b/src/websocket/listener/commands/chatgpt-usage.ts @@ -4,7 +4,7 @@ import type { ChatGPTUsageReadCommand, ChatGPTUsageReadResponseMessage, } from "@/types/protocol_v2"; -import { isChatGPTUsageReadCommand } from "@/websocket/listener/protocol-inbound"; +import { isChatGPTUsageReadCommand } from "@/websocket/listener/chatgpt-usage-protocol-inbound"; import type { RunDetachedListenerTask, SafeSocketSend } from "./types"; type ChatGPTUsageCommandContext = { diff --git a/src/websocket/listener/message-router.ts b/src/websocket/listener/message-router.ts index c82184a639..f2e4dee90d 100644 --- a/src/websocket/listener/message-router.ts +++ b/src/websocket/listener/message-router.ts @@ -21,6 +21,7 @@ import { import { handleExecuteCommand } from "./commands"; import { handleAgentConversationManagementProtocolCommand } from "./commands/agents-conversations"; import { handleAppServerInfoCommand } from "./commands/app-server-info"; +import { handleChatGPTResetCreditsCommand } from "./commands/chatgpt-reset-credits"; import { handleChatGPTUsageCommand } from "./commands/chatgpt-usage"; import { handleConnectProvidersCommand } from "./commands/connect-providers"; import { handleCronProtocolCommand } from "./commands/cron"; @@ -673,6 +674,16 @@ export function createListenerMessageHandler( return; } + if ( + handleChatGPTResetCreditsCommand(parsed, { + socket, + safeSocketSend, + runDetachedListenerTask, + }) + ) { + return; + } + if ( handleCronProtocolCommand(parsed, { socket, diff --git a/src/websocket/listener/protocol-inbound.ts b/src/websocket/listener/protocol-inbound.ts index 5d19731535..95bd95efe5 100644 --- a/src/websocket/listener/protocol-inbound.ts +++ b/src/websocket/listener/protocol-inbound.ts @@ -37,7 +37,6 @@ import type { ChannelsListCommand, ChannelTargetBindCommand, ChannelTargetsListCommand, - ChatGPTUsageReadCommand, CheckoutBranchCommand, ClientToolsetConfig, ConnectProviderCommand, @@ -103,6 +102,17 @@ import type { WriteMemoryFileCommand, WsProtocolCommand, } from "@/types/protocol_v2"; +import { + isChatGPTRateLimitResetCreditConsumeCommand, + isChatGPTRateLimitResetCreditsListCommand, + isChatGPTUsageReadCommand, +} from "@/websocket/listener/chatgpt-usage-protocol-inbound"; + +export { + isChatGPTRateLimitResetCreditConsumeCommand, + isChatGPTRateLimitResetCreditsListCommand, + isChatGPTUsageReadCommand, +} from "@/websocket/listener/chatgpt-usage-protocol-inbound"; const EXPERIMENT_IDS = new Set([ "conversation_titles", @@ -994,26 +1004,6 @@ export function isDisconnectProviderCommand( ); } -export function isChatGPTUsageReadCommand( - value: unknown, -): value is ChatGPTUsageReadCommand { - if (!value || typeof value !== "object") return false; - const c = value as { - type?: unknown; - request_id?: unknown; - target?: unknown; - provider_name?: unknown; - force_refresh?: unknown; - }; - return ( - c.type === "chatgpt_usage_read" && - typeof c.request_id === "string" && - (c.target === "local" || c.target === "api") && - (c.provider_name === undefined || typeof c.provider_name === "string") && - (c.force_refresh === undefined || typeof c.force_refresh === "boolean") - ); -} - export function isUpdateModelCommand( value: unknown, ): value is UpdateModelCommand { @@ -2187,6 +2177,8 @@ export function parseServerMessage( isConnectProviderCommand(parsed) || isDisconnectProviderCommand(parsed) || isChatGPTUsageReadCommand(parsed) || + isChatGPTRateLimitResetCreditsListCommand(parsed) || + isChatGPTRateLimitResetCreditConsumeCommand(parsed) || isUpdateModelCommand(parsed) || isUpdateToolsetCommand(parsed) || isCronListCommand(parsed) ||