Add Codex usage indicator

Author: lintowe

apps/desktop/src/clientPersistence.test.ts

@@ -50,6 +50,7 @@ function makeSecretStorage(available: boolean): DesktopSecretStorage {
 
 const clientSettings: ClientSettings = {
   autoOpenPlanSidebar: false,
+  codexUsageIndicatorMode: "five-hour",
   confirmThreadArchive: true,
   confirmThreadDelete: false,
   diffIgnoreWhitespace: true,

apps/server/src/orchestration/Layers/CheckpointReactor.test.ts

@@ -114,6 +114,7 @@ function createProviderServiceHarness(
           continuationKey: `${providerName}:instance:${instanceId}`,
         },
       }),
+    getCodexUsage: () => Effect.succeed(null),
     rollbackConversation,
     get streamEvents() {
       return Stream.fromPubSub(runtimeEventPubSub);

apps/server/src/orchestration/Layers/ProviderCommandReactor.test.ts

@@ -304,6 +304,7 @@ describe("ProviderCommandReactor", () => {
           },
         });
       },
+      getCodexUsage: () => Effect.succeed(null),
       rollbackConversation: () => unsupported(),
       get streamEvents() {
         return Stream.fromPubSub(runtimeEventPubSub);

apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts

@@ -113,6 +113,7 @@ function createProviderServiceHarness() {
         },
       });
     },
+    getCodexUsage: () => Effect.succeed(null),
     rollbackConversation: () => unsupported(),
     get streamEvents() {
       return Stream.fromPubSub(runtimeEventPubSub);

apps/server/src/provider/Layers/CodexAdapter.test.ts

@@ -23,6 +23,7 @@ import { it, vi } from "@effect/vitest";
 
 import { Context, Effect, Exit, Fiber, Layer, Option, Queue, Schema, Scope, Stream } from "effect";
 import * as CodexErrors from "effect-codex-app-server/errors";
+import type * as EffectCodexSchema from "effect-codex-app-server/schema";
 
 import { ServerConfig } from "../../config.ts";
 import { ServerSettingsService } from "../../serverSettings.ts";
@@ -92,6 +93,15 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape {
       }),
   );
 
+  public readonly readAccountRateLimitsImpl = vi.fn(
+    (): Promise<EffectCodexSchema.V2GetAccountRateLimitsResponse> =>
+      Promise.resolve({
+        rateLimits: {
+          primary: { usedPercent: 25, windowDurationMins: 300 },
+        },
+      }),
+  );
+
   public readonly respondToRequestImpl = vi.fn(
     (_requestId: ApprovalRequestId, _decision: ProviderApprovalDecision): Promise<void> =>
       Promise.resolve(undefined),
@@ -130,6 +140,8 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape {
     return Effect.promise(() => this.rollbackThreadImpl(numTurns));
   }
 
+  readAccountRateLimits = Effect.promise(() => this.readAccountRateLimitsImpl());
+
   respondToRequest(requestId: ApprovalRequestId, decision: ProviderApprovalDecision) {
     return Effect.promise(() => this.respondToRequestImpl(requestId, decision));
   }
@@ -159,6 +171,7 @@ function makeRuntimeFactory() {
 
   return {
     factory,
+    runtimes,
     get lastRuntime(): FakeCodexRuntime | undefined {
       return runtimes.at(-1);
     },
@@ -348,6 +361,95 @@ sessionErrorLayer("CodexAdapterLive session errors", (it) => {
     }),
   );
 
+  it.effect("reads and normalizes account rate limits through the active runtime", () =>
+    Effect.gen(function* () {
+      const adapter = yield* CodexAdapter;
+      yield* adapter.startSession({
+        provider: ProviderDriverKind.make("codex"),
+        threadId: asThreadId("usage-thread"),
+        runtimeMode: "full-access",
+      });
+      const runtime = sessionRuntimeFactory.lastRuntime;
+      assert.ok(runtime);
+      runtime.readAccountRateLimitsImpl.mockResolvedValueOnce({
+        rateLimits: {
+          primary: { usedPercent: 30, windowDurationMins: 300 },
+          secondary: { usedPercent: 80, windowDurationMins: 10_080 },
+        },
+      });
+
+      const snapshot = yield* adapter.readCodexUsage!();
+
+      assert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, 1);
+      assert.deepStrictEqual(
+        snapshot?.windows.map((window) => ({
+          kind: window.kind,
+          remainingPercent: window.remainingPercent,
+        })),
+        [
+          { kind: "five-hour", remainingPercent: 70 },
+          { kind: "weekly", remainingPercent: 20 },
+        ],
+      );
+    }),
+  );
+
+  it.effect("keeps cached account rate limits when an active read has no displayable windows", () =>
+    Effect.gen(function* () {
+      const adapter = yield* CodexAdapter;
+      yield* adapter.startSession({
+        provider: ProviderDriverKind.make("codex"),
+        threadId: asThreadId("usage-cache-thread"),
+        runtimeMode: "full-access",
+      });
+      const runtime = sessionRuntimeFactory.lastRuntime;
+      assert.ok(runtime);
+      runtime.readAccountRateLimitsImpl.mockResolvedValueOnce({
+        rateLimits: {
+          primary: { usedPercent: 45, windowDurationMins: 300 },
+        },
+      });
+      yield* adapter.readCodexUsage!();
+      runtime.readAccountRateLimitsImpl.mockResolvedValueOnce({
+        rateLimits: {},
+      });
+
+      const snapshot = yield* adapter.readCodexUsage!();
+
+      assert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, 2);
+      assert.equal(snapshot?.source, "cache");
+      assert.deepStrictEqual(snapshot?.windows[0], {
+        kind: "five-hour",
+        usedPercent: 45,
+        remainingPercent: 55,
+        resetsAt: null,
+        windowDurationMins: 300,
+      });
+    }),
+  );
+
+  it.effect("reads account rate limits even before a Codex thread session exists", () =>
+    Effect.gen(function* () {
+      const adapter = yield* CodexAdapter;
+      yield* adapter.stopAll();
+      const snapshot = yield* adapter.readCodexUsage!();
+      const runtime = sessionRuntimeFactory.lastRuntime;
+
+      assert.ok(runtime);
+      assert.equal(runtime.options.threadId, asThreadId("codex-usage"));
+      assert.equal(runtime.startImpl.mock.calls.length, 1);
+      assert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, 1);
+      assert.equal(runtime.closeImpl.mock.calls.length, 1);
+      assert.deepStrictEqual(snapshot?.windows[0], {
+        kind: "five-hour",
+        usedPercent: 25,
+        remainingPercent: 75,
+        resetsAt: null,
+        windowDurationMins: 300,
+      });
+    }),
+  );
+
   it.effect("maps codex model options for the adapter's bound custom instance id", () => {
     const customInstanceId = ProviderInstanceId.make("codex_personal");
     const customRuntimeFactory = makeRuntimeFactory();

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

@@ -17,6 +17,7 @@ import {
   type ProviderRuntimeEvent,
   type ProviderRequestKind,
   type ThreadTokenUsageSnapshot,
+  type CodexUsageSnapshot,
   type ProviderUserInputAnswers,
   RuntimeItemId,
   RuntimeRequestId,
@@ -54,6 +55,7 @@ import {
   type CodexSessionRuntimeShape,
 } from "./CodexSessionRuntime.ts";
 import { type EventNdjsonLogger, makeEventNdjsonLogger } from "./EventNdjsonLogger.ts";
+import { normalizeCodexUsageSnapshot } from "../codexUsage.ts";
 
 const PROVIDER = ProviderDriverKind.make("codex");
 
@@ -1350,6 +1352,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
     options?.nativeEventLogger === undefined ? nativeEventLogger : undefined;
   const runtimeEventQueue = yield* Queue.unbounded<ProviderRuntimeEvent>();
   const sessions = new Map<ThreadId, CodexAdapterSessionContext>();
+  let cachedCodexUsage: CodexUsageSnapshot | null = null;
 
   const startSession: CodexAdapterShape["startSession"] = (input) =>
     Effect.scoped(
@@ -1409,6 +1412,19 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
         const eventFiber = yield* Stream.runForEach(runtime.events, (event) =>
           Effect.gen(function* () {
             yield* writeNativeEvent(event);
+            if (event.method === "account/rateLimits/updated") {
+              const payload = readPayload(
+                EffectCodexSchema.V2AccountRateLimitsUpdatedNotification,
+                event.payload,
+              );
+              if (payload) {
+                cachedCodexUsage = normalizeCodexUsageSnapshot({
+                  providerInstanceId: boundInstanceId,
+                  payload,
+                  source: "notification",
+                });
+              }
+            }
             const runtimeEvents = mapToRuntimeEvents(event, event.threadId);
             if (runtimeEvents.length === 0) {
               yield* Effect.logDebug("ignoring unhandled Codex provider event", {
@@ -1644,6 +1660,90 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
   const hasSession: CodexAdapterShape["hasSession"] = (threadId) =>
     Effect.succeed(Boolean(sessions.get(threadId) && !sessions.get(threadId)?.stopped));
 
+  const readCodexUsageWithoutSession = Effect.fn("readCodexUsageWithoutSession")(function* () {
+    const usageThreadId = ThreadId.make("codex-usage");
+    const createRuntime = options?.makeRuntime ?? makeCodexSessionRuntime;
+    return yield* Effect.acquireUseRelease(
+      Scope.make("sequential"),
+      (usageScope) =>
+        Effect.gen(function* () {
+          const runtime = yield* createRuntime({
+            threadId: usageThreadId,
+            providerInstanceId: boundInstanceId,
+            cwd: process.cwd(),
+            binaryPath: codexConfig.binaryPath,
+            ...(options?.environment ? { environment: options.environment } : {}),
+            ...(codexConfig.homePath ? { homePath: codexConfig.homePath } : {}),
+            runtimeMode: "full-access",
+          }).pipe(
+            Effect.provideService(Scope.Scope, usageScope),
+            Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, childProcessSpawner),
+            Effect.mapError(
+              (cause) =>
+                new ProviderAdapterProcessError({
+                  provider: PROVIDER,
+                  threadId: usageThreadId,
+                  detail: cause.message,
+                  cause,
+                }),
+            ),
+          );
+          const payload = yield* runtime.start().pipe(
+            Effect.mapError(
+              (cause) =>
+                new ProviderAdapterProcessError({
+                  provider: PROVIDER,
+                  threadId: usageThreadId,
+                  detail: cause.message,
+                  cause,
+                }),
+            ),
+            Effect.andThen(
+              runtime.readAccountRateLimits.pipe(
+                Effect.mapError((cause) =>
+                  mapCodexRuntimeError(usageThreadId, "account/rateLimits/read", cause),
+                ),
+              ),
+            ),
+            Effect.ensuring(runtime.close),
+          );
+          return normalizeCodexUsageSnapshot({
+            providerInstanceId: boundInstanceId,
+            payload,
+            source: "read",
+          });
+        }),
+      (usageScope) => Scope.close(usageScope, Exit.void),
+    );
+  });
+
+  const readCodexUsage: CodexAdapterShape["readCodexUsage"] = Effect.fn("readCodexUsage")(
+    function* () {
+      const session = Array.from(sessions.values()).findLast((candidate) => !candidate.stopped);
+      if (!session) {
+        const snapshot = yield* readCodexUsageWithoutSession();
+        cachedCodexUsage = snapshot ?? cachedCodexUsage;
+        return (
+          snapshot ?? (cachedCodexUsage ? { ...cachedCodexUsage, source: "cache" as const } : null)
+        );
+      }
+      const payload = yield* session.runtime.readAccountRateLimits.pipe(
+        Effect.mapError((cause) =>
+          mapCodexRuntimeError(session.threadId, "account/rateLimits/read", cause),
+        ),
+      );
+      const snapshot = normalizeCodexUsageSnapshot({
+        providerInstanceId: boundInstanceId,
+        payload,
+        source: "read",
+      });
+      cachedCodexUsage = snapshot ?? cachedCodexUsage;
+      return (
+        snapshot ?? (cachedCodexUsage ? { ...cachedCodexUsage, source: "cache" as const } : null)
+      );
+    },
+  );
+
   const stopAll: CodexAdapterShape["stopAll"] = () =>
     Effect.forEach(Array.from(sessions.values()), stopSessionInternal, {
       concurrency: 1,
@@ -1673,6 +1773,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
     stopSession,
     listSessions,
     hasSession,
+    readCodexUsage,
     stopAll,
     get streamEvents() {
       return Stream.fromQueue(runtimeEventQueue);

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

@@ -120,6 +120,10 @@ export interface CodexSessionRuntimeShape {
   readonly rollbackThread: (
     numTurns: number,
   ) => Effect.Effect<CodexThreadSnapshot, CodexSessionRuntimeError>;
+  readonly readAccountRateLimits: Effect.Effect<
+    EffectCodexSchema.V2GetAccountRateLimitsResponse,
+    CodexSessionRuntimeError
+  >;
   readonly respondToRequest: (
     requestId: ApprovalRequestId,
     decision: ProviderApprovalDecision,
@@ -1286,6 +1290,7 @@ export const makeCodexSessionRuntime = (
           });
           return parseThreadSnapshot(response);
         }),
+      readAccountRateLimits: client.request("account/rateLimits/read", undefined),
       respondToRequest: (requestId, decision) =>
         Effect.gen(function* () {
           const pending = (yield* Ref.get(pendingApprovalsRef)).get(requestId);