feat: add account rate limits handling and corresponding tests

Author: Aditya190803

apps/server/src/codexAppServerManager.ts

@@ -519,6 +519,25 @@ export class CodexAppServerManager extends EventEmitter<CodexAppServerManagerEve
         console.log("codex account/read failed", error);
       }
 
+      try {
+        const accountRateLimitsResponse = await this.sendRequest(
+          context,
+          "account/rateLimits/read",
+          {},
+        );
+        this.emitEvent({
+          id: EventId.makeUnsafe(randomUUID()),
+          kind: "notification",
+          provider: "codex",
+          threadId: context.session.threadId,
+          createdAt: new Date().toISOString(),
+          method: "account/rateLimits/updated",
+          payload: accountRateLimitsResponse,
+        });
+      } catch (error) {
+        console.log("codex account/rateLimits/read failed", error);
+      }
+
       const normalizedModel = resolveCodexModelForAccount(
         normalizeCodexModelSlug(input.model),
         context.account,

apps/server/src/provider/Layers/CodexProvider.ts

@@ -7,18 +7,7 @@ import type {
   ServerProviderAuth,
   ServerProviderState,
 } from "@t3tools/contracts";
-import {
-  Cache,
-  Duration,
-  Effect,
-  Equal,
-  FileSystem,
-  Layer,
-  Option,
-  Path,
-  Result,
-  Stream,
-} from "effect";
+import { Effect, Equal, FileSystem, Layer, Option, Path, Result, Stream } from "effect";
 import { ChildProcess, ChildProcessSpawner } from "effect/unstable/process";
 
 import {
@@ -562,20 +551,11 @@ export const CodexProviderLive = Layer.effect(
     const fileSystem = yield* FileSystem.FileSystem;
     const path = yield* Path.Path;
     const spawner = yield* ChildProcessSpawner.ChildProcessSpawner;
-    const accountProbeCache = yield* Cache.make({
-      capacity: 4,
-      timeToLive: Duration.minutes(5),
-      lookup: (key: string) => {
-        const [binaryPath, homePath] = JSON.parse(key) as [string, string | undefined];
-        return probeCodexCapabilities({
-          binaryPath,
-          ...(homePath ? { homePath } : {}),
-        });
-      },
-    });
-
     const checkProvider = checkCodexProviderStatus((input) =>
-      Cache.get(accountProbeCache, JSON.stringify([input.binaryPath, input.homePath])),
+      probeCodexCapabilities({
+        binaryPath: input.binaryPath,
+        ...(input.homePath ? { homePath: input.homePath } : {}),
+      }),
     ).pipe(
       Effect.provideService(ServerSettingsService, serverSettings),
       Effect.provideService(FileSystem.FileSystem, fileSystem),

apps/server/src/provider/codexAppServer.test.ts

@@ -0,0 +1,116 @@
+import { chmodSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
+import os from "node:os";
+import path from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+
+import { probeCodexAccountState } from "./codexAppServer";
+
+const tempDirs: Array<string> = [];
+
+afterEach(() => {
+  for (const tempDir of tempDirs.splice(0)) {
+    rmSync(tempDir, { recursive: true, force: true });
+  }
+});
+
+function createCodexProbeStub(options: {
+  readonly rateLimitsBehavior: "ignore" | "respond";
+}): string {
+  const tempDir = mkdtempSync(path.join(os.tmpdir(), "codex-probe-stub-"));
+  tempDirs.push(tempDir);
+
+  const scriptPath = path.join(tempDir, "codex-stub.mjs");
+  const respondsRateLimits = options.rateLimitsBehavior === "respond";
+  const content = `#!/usr/bin/env node
+import readline from "node:readline";
+
+const output = (message) => {
+  process.stdout.write(JSON.stringify(message) + "\\n");
+};
+
+const reader = readline.createInterface({ input: process.stdin });
+reader.on("line", (line) => {
+  const message = JSON.parse(line);
+
+  if (message.method === "initialize") {
+    output({ id: 1, result: {} });
+    return;
+  }
+
+  if (message.id === 2 && message.method === "account/read") {
+    output({
+      id: 2,
+      result: {
+        account: {
+          type: "chatgpt",
+          planType: "pro",
+        },
+      },
+    });
+    return;
+  }
+
+  if (message.id === 3 && message.method === "account/rateLimits/read") {
+    if (${respondsRateLimits}) {
+      output({
+        id: 3,
+        result: {
+          rateLimits: {
+            primary: {
+              remaining: 7,
+              used: 3,
+            },
+          },
+        },
+      });
+    }
+  }
+});
+`;
+  writeFileSync(scriptPath, content, { encoding: "utf8" });
+  chmodSync(scriptPath, 0o755);
+  return scriptPath;
+}
+
+describe("probeCodexAccountState", () => {
+  it("resolves when account/rateLimits/read is ignored", async () => {
+    const binaryPath = createCodexProbeStub({ rateLimitsBehavior: "ignore" });
+
+    const state = await probeCodexAccountState({
+      binaryPath,
+      signal: AbortSignal.timeout(1_000),
+    });
+
+    expect(state.snapshot).toEqual({
+      type: "chatgpt",
+      planType: "pro",
+      sparkEnabled: true,
+    });
+    expect(state.account).toEqual({
+      type: "chatgpt",
+      planType: "pro",
+    });
+    expect(state.rateLimits).toBeNull();
+  });
+
+  it("includes rate limits when account/rateLimits/read responds", async () => {
+    const binaryPath = createCodexProbeStub({ rateLimitsBehavior: "respond" });
+
+    const state = await probeCodexAccountState({
+      binaryPath,
+      signal: AbortSignal.timeout(1_000),
+    });
+
+    expect(state.snapshot).toEqual({
+      type: "chatgpt",
+      planType: "pro",
+      sparkEnabled: true,
+    });
+    expect(state.rateLimits).toEqual({
+      primary: {
+        remaining: 7,
+        used: 3,
+      },
+    });
+  });
+});

apps/server/src/provider/codexAppServer.ts

@@ -15,12 +15,16 @@ export interface CodexAccountState {
 
 interface JsonRpcProbeResponse {
   readonly id?: unknown;
+  readonly method?: unknown;
+  readonly params?: unknown;
   readonly result?: unknown;
   readonly error?: {
     readonly message?: unknown;
   };
 }
 
+const RATE_LIMITS_PROBE_GRACE_MS = 1_500;
+
 function readErrorMessage(response: JsonRpcProbeResponse): string | undefined {
   return typeof response.error?.message === "string" ? response.error.message : undefined;
 }
@@ -80,8 +84,13 @@ export async function probeCodexAccountState(input: {
     let accountSnapshot: CodexAccountSnapshot | null = null;
     let accountPayload: Record<string, unknown> | null = null;
     let rateLimitsPayload: Record<string, unknown> | null | undefined;
+    let rateLimitsFallbackTimer: NodeJS.Timeout | undefined;
 
     const cleanup = () => {
+      if (rateLimitsFallbackTimer) {
+        clearTimeout(rateLimitsFallbackTimer);
+        rateLimitsFallbackTimer = undefined;
+      }
       output.removeAllListeners();
       output.close();
       child.removeAllListeners();
@@ -107,7 +116,17 @@ export async function probeCodexAccountState(input: {
       );
 
     const maybeFinish = () => {
-      if (!accountSnapshot || rateLimitsPayload === undefined) {
+      if (!accountSnapshot) {
+        return;
+      }
+
+      if (rateLimitsPayload === undefined) {
+        if (!rateLimitsFallbackTimer) {
+          rateLimitsFallbackTimer = setTimeout(() => {
+            rateLimitsPayload = null;
+            maybeFinish();
+          }, RATE_LIMITS_PROBE_GRACE_MS);
+        }
         return;
       }
 
@@ -153,6 +172,15 @@ export async function probeCodexAccountState(input: {
       }
 
       const response = parsed as JsonRpcProbeResponse;
+      if (response.method === "account/rateLimits/updated") {
+        const payload =
+          readCodexRateLimitsPayload(response.params) ??
+          readCodexRateLimitsPayload(response.result);
+        rateLimitsPayload = payload ?? null;
+        maybeFinish();
+        return;
+      }
+
       if (response.id === 1) {
         const errorMessage = readErrorMessage(response);
         if (errorMessage) {