Skip to content

Commit 4e9f527

Browse files
committed
Add Codex usage indicator
1 parent f7748a0 commit 4e9f527

30 files changed

Lines changed: 1092 additions & 5 deletions

apps/desktop/src/clientPersistence.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ function makeSecretStorage(available: boolean): DesktopSecretStorage {
5050

5151
const clientSettings: ClientSettings = {
5252
autoOpenPlanSidebar: false,
53+
codexUsageIndicatorMode: "five-hour",
5354
confirmThreadArchive: true,
5455
confirmThreadDelete: false,
5556
diffIgnoreWhitespace: true,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ function createProviderServiceHarness(
114114
continuationKey: `${providerName}:instance:${instanceId}`,
115115
},
116116
}),
117+
getCodexUsage: () => Effect.succeed(null),
117118
rollbackConversation,
118119
get streamEvents() {
119120
return Stream.fromPubSub(runtimeEventPubSub);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,7 @@ describe("ProviderCommandReactor", () => {
304304
},
305305
});
306306
},
307+
getCodexUsage: () => Effect.succeed(null),
307308
rollbackConversation: () => unsupported(),
308309
get streamEvents() {
309310
return Stream.fromPubSub(runtimeEventPubSub);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ function createProviderServiceHarness() {
113113
},
114114
});
115115
},
116+
getCodexUsage: () => Effect.succeed(null),
116117
rollbackConversation: () => unsupported(),
117118
get streamEvents() {
118119
return Stream.fromPubSub(runtimeEventPubSub);

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import { it, vi } from "@effect/vitest";
2323

2424
import { Context, Effect, Exit, Fiber, Layer, Option, Queue, Schema, Scope, Stream } from "effect";
2525
import * as CodexErrors from "effect-codex-app-server/errors";
26+
import type * as EffectCodexSchema from "effect-codex-app-server/schema";
2627

2728
import { ServerConfig } from "../../config.ts";
2829
import { ServerSettingsService } from "../../serverSettings.ts";
@@ -92,6 +93,15 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape {
9293
}),
9394
);
9495

96+
public readonly readAccountRateLimitsImpl = vi.fn(
97+
(): Promise<EffectCodexSchema.V2GetAccountRateLimitsResponse> =>
98+
Promise.resolve({
99+
rateLimits: {
100+
primary: { usedPercent: 25, windowDurationMins: 300 },
101+
},
102+
}),
103+
);
104+
95105
public readonly respondToRequestImpl = vi.fn(
96106
(_requestId: ApprovalRequestId, _decision: ProviderApprovalDecision): Promise<void> =>
97107
Promise.resolve(undefined),
@@ -130,6 +140,8 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape {
130140
return Effect.promise(() => this.rollbackThreadImpl(numTurns));
131141
}
132142

143+
readAccountRateLimits = Effect.promise(() => this.readAccountRateLimitsImpl());
144+
133145
respondToRequest(requestId: ApprovalRequestId, decision: ProviderApprovalDecision) {
134146
return Effect.promise(() => this.respondToRequestImpl(requestId, decision));
135147
}
@@ -159,6 +171,7 @@ function makeRuntimeFactory() {
159171

160172
return {
161173
factory,
174+
runtimes,
162175
get lastRuntime(): FakeCodexRuntime | undefined {
163176
return runtimes.at(-1);
164177
},
@@ -348,6 +361,60 @@ sessionErrorLayer("CodexAdapterLive session errors", (it) => {
348361
}),
349362
);
350363

364+
it.effect("reads and normalizes account rate limits through the active runtime", () =>
365+
Effect.gen(function* () {
366+
const adapter = yield* CodexAdapter;
367+
yield* adapter.startSession({
368+
provider: ProviderDriverKind.make("codex"),
369+
threadId: asThreadId("usage-thread"),
370+
runtimeMode: "full-access",
371+
});
372+
const runtime = sessionRuntimeFactory.lastRuntime;
373+
assert.ok(runtime);
374+
runtime.readAccountRateLimitsImpl.mockResolvedValueOnce({
375+
rateLimits: {
376+
primary: { usedPercent: 30, windowDurationMins: 300 },
377+
secondary: { usedPercent: 80, windowDurationMins: 10_080 },
378+
},
379+
});
380+
381+
const snapshot = yield* adapter.readCodexUsage!();
382+
383+
assert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, 1);
384+
assert.deepStrictEqual(
385+
snapshot?.windows.map((window) => ({
386+
kind: window.kind,
387+
remainingPercent: window.remainingPercent,
388+
})),
389+
[
390+
{ kind: "five-hour", remainingPercent: 70 },
391+
{ kind: "weekly", remainingPercent: 20 },
392+
],
393+
);
394+
}),
395+
);
396+
397+
it.effect("reads account rate limits even before a Codex thread session exists", () =>
398+
Effect.gen(function* () {
399+
const adapter = yield* CodexAdapter;
400+
yield* adapter.stopAll();
401+
const snapshot = yield* adapter.readCodexUsage!();
402+
const runtime = sessionRuntimeFactory.lastRuntime;
403+
404+
assert.ok(runtime);
405+
assert.equal(runtime.options.threadId, asThreadId("codex-usage"));
406+
assert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, 1);
407+
assert.equal(runtime.closeImpl.mock.calls.length, 1);
408+
assert.deepStrictEqual(snapshot?.windows[0], {
409+
kind: "five-hour",
410+
usedPercent: 25,
411+
remainingPercent: 75,
412+
resetsAt: null,
413+
windowDurationMins: 300,
414+
});
415+
}),
416+
);
417+
351418
it.effect("maps codex model options for the adapter's bound custom instance id", () => {
352419
const customInstanceId = ProviderInstanceId.make("codex_personal");
353420
const customRuntimeFactory = makeRuntimeFactory();

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

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
type ProviderRuntimeEvent,
1818
type ProviderRequestKind,
1919
type ThreadTokenUsageSnapshot,
20+
type CodexUsageSnapshot,
2021
type ProviderUserInputAnswers,
2122
RuntimeItemId,
2223
RuntimeRequestId,
@@ -54,6 +55,7 @@ import {
5455
type CodexSessionRuntimeShape,
5556
} from "./CodexSessionRuntime.ts";
5657
import { type EventNdjsonLogger, makeEventNdjsonLogger } from "./EventNdjsonLogger.ts";
58+
import { normalizeCodexUsageSnapshot } from "../codexUsage.ts";
5759

5860
const PROVIDER = ProviderDriverKind.make("codex");
5961

@@ -1350,6 +1352,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
13501352
options?.nativeEventLogger === undefined ? nativeEventLogger : undefined;
13511353
const runtimeEventQueue = yield* Queue.unbounded<ProviderRuntimeEvent>();
13521354
const sessions = new Map<ThreadId, CodexAdapterSessionContext>();
1355+
let cachedCodexUsage: CodexUsageSnapshot | null = null;
13531356

13541357
const startSession: CodexAdapterShape["startSession"] = (input) =>
13551358
Effect.scoped(
@@ -1409,6 +1412,19 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
14091412
const eventFiber = yield* Stream.runForEach(runtime.events, (event) =>
14101413
Effect.gen(function* () {
14111414
yield* writeNativeEvent(event);
1415+
if (event.method === "account/rateLimits/updated") {
1416+
const payload = readPayload(
1417+
EffectCodexSchema.V2AccountRateLimitsUpdatedNotification,
1418+
event.payload,
1419+
);
1420+
if (payload) {
1421+
cachedCodexUsage = normalizeCodexUsageSnapshot({
1422+
providerInstanceId: boundInstanceId,
1423+
payload,
1424+
source: "notification",
1425+
});
1426+
}
1427+
}
14121428
const runtimeEvents = mapToRuntimeEvents(event, event.threadId);
14131429
if (runtimeEvents.length === 0) {
14141430
yield* Effect.logDebug("ignoring unhandled Codex provider event", {
@@ -1644,6 +1660,75 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
16441660
const hasSession: CodexAdapterShape["hasSession"] = (threadId) =>
16451661
Effect.succeed(Boolean(sessions.get(threadId) && !sessions.get(threadId)?.stopped));
16461662

1663+
const readCodexUsageWithoutSession = Effect.fn("readCodexUsageWithoutSession")(function* () {
1664+
const usageThreadId = ThreadId.make("codex-usage");
1665+
const createRuntime = options?.makeRuntime ?? makeCodexSessionRuntime;
1666+
return yield* Effect.acquireUseRelease(
1667+
Scope.make("sequential"),
1668+
(usageScope) =>
1669+
Effect.gen(function* () {
1670+
const runtime = yield* createRuntime({
1671+
threadId: usageThreadId,
1672+
providerInstanceId: boundInstanceId,
1673+
cwd: process.cwd(),
1674+
binaryPath: codexConfig.binaryPath,
1675+
...(options?.environment ? { environment: options.environment } : {}),
1676+
...(codexConfig.homePath ? { homePath: codexConfig.homePath } : {}),
1677+
runtimeMode: "full-access",
1678+
}).pipe(
1679+
Effect.provideService(Scope.Scope, usageScope),
1680+
Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, childProcessSpawner),
1681+
Effect.mapError(
1682+
(cause) =>
1683+
new ProviderAdapterProcessError({
1684+
provider: PROVIDER,
1685+
threadId: usageThreadId,
1686+
detail: cause.message,
1687+
cause,
1688+
}),
1689+
),
1690+
);
1691+
const payload = yield* runtime.readAccountRateLimits.pipe(
1692+
Effect.mapError((cause) =>
1693+
mapCodexRuntimeError(usageThreadId, "account/rateLimits/read", cause),
1694+
),
1695+
Effect.ensuring(runtime.close),
1696+
);
1697+
return normalizeCodexUsageSnapshot({
1698+
providerInstanceId: boundInstanceId,
1699+
payload,
1700+
source: "read",
1701+
});
1702+
}),
1703+
(usageScope) => Scope.close(usageScope, Exit.void),
1704+
);
1705+
});
1706+
1707+
const readCodexUsage: CodexAdapterShape["readCodexUsage"] = Effect.fn("readCodexUsage")(
1708+
function* () {
1709+
const session = Array.from(sessions.values()).findLast((candidate) => !candidate.stopped);
1710+
if (!session) {
1711+
const snapshot = yield* readCodexUsageWithoutSession();
1712+
cachedCodexUsage = snapshot ?? cachedCodexUsage;
1713+
return (
1714+
snapshot ?? (cachedCodexUsage ? { ...cachedCodexUsage, source: "cache" as const } : null)
1715+
);
1716+
}
1717+
const payload = yield* session.runtime.readAccountRateLimits.pipe(
1718+
Effect.mapError((cause) =>
1719+
mapCodexRuntimeError(session.threadId, "account/rateLimits/read", cause),
1720+
),
1721+
);
1722+
const snapshot = normalizeCodexUsageSnapshot({
1723+
providerInstanceId: boundInstanceId,
1724+
payload,
1725+
source: "read",
1726+
});
1727+
cachedCodexUsage = snapshot;
1728+
return snapshot;
1729+
},
1730+
);
1731+
16471732
const stopAll: CodexAdapterShape["stopAll"] = () =>
16481733
Effect.forEach(Array.from(sessions.values()), stopSessionInternal, {
16491734
concurrency: 1,
@@ -1673,6 +1758,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* (
16731758
stopSession,
16741759
listSessions,
16751760
hasSession,
1761+
readCodexUsage,
16761762
stopAll,
16771763
get streamEvents() {
16781764
return Stream.fromQueue(runtimeEventQueue);

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,10 @@ export interface CodexSessionRuntimeShape {
120120
readonly rollbackThread: (
121121
numTurns: number,
122122
) => Effect.Effect<CodexThreadSnapshot, CodexSessionRuntimeError>;
123+
readonly readAccountRateLimits: Effect.Effect<
124+
EffectCodexSchema.V2GetAccountRateLimitsResponse,
125+
CodexSessionRuntimeError
126+
>;
123127
readonly respondToRequest: (
124128
requestId: ApprovalRequestId,
125129
decision: ProviderApprovalDecision,
@@ -1286,6 +1290,7 @@ export const makeCodexSessionRuntime = (
12861290
});
12871291
return parseThreadSnapshot(response);
12881292
}),
1293+
readAccountRateLimits: client.request("account/rateLimits/read", undefined),
12891294
respondToRequest: (requestId, decision) =>
12901295
Effect.gen(function* () {
12911296
const pending = (yield* Ref.get(pendingApprovalsRef)).get(requestId);

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import type {
88
ProviderSendTurnInput,
99
ProviderSession,
1010
ProviderTurnStartResult,
11+
CodexUsageSnapshot,
1112
} from "@t3tools/contracts";
1213
import {
1314
ApprovalRequestId,
@@ -191,6 +192,24 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) {
191192
sessions.clear();
192193
}),
193194
);
195+
const readCodexUsage = vi.fn(
196+
(): Effect.Effect<CodexUsageSnapshot | null, ProviderAdapterError> =>
197+
Effect.succeed({
198+
providerInstanceId: codexInstanceId,
199+
checkedAt: "2026-05-04T00:00:00.000Z",
200+
windows: [
201+
{
202+
kind: "five-hour",
203+
usedPercent: 25,
204+
remainingPercent: 75,
205+
resetsAt: null,
206+
windowDurationMins: 300,
207+
},
208+
],
209+
rateLimitReachedType: null,
210+
source: "read",
211+
}),
212+
);
194213

195214
const adapter: ProviderAdapterShape<ProviderAdapterError> = {
196215
provider,
@@ -207,6 +226,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) {
207226
hasSession,
208227
readThread,
209228
rollbackThread,
229+
...(provider === CODEX_DRIVER ? { readCodexUsage } : {}),
210230
stopAll,
211231
get streamEvents() {
212232
return Stream.fromPubSub(runtimeEventPubSub);
@@ -243,6 +263,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) {
243263
readThread,
244264
rollbackThread,
245265
stopAll,
266+
readCodexUsage,
246267
};
247268
}
248269

@@ -772,6 +793,20 @@ it.effect(
772793
);
773794

774795
routing.layer("ProviderServiceLive routing", (it) => {
796+
it.effect("returns usage for Codex instances and null for non-Codex instances", () =>
797+
Effect.gen(function* () {
798+
const provider = yield* ProviderService;
799+
800+
const codexUsage = yield* provider.getCodexUsage(codexInstanceId);
801+
const claudeUsage = yield* provider.getCodexUsage(claudeAgentInstanceId);
802+
803+
assert.equal(codexUsage?.windows[0]?.remainingPercent, 75);
804+
assert.equal(routing.codex.readCodexUsage.mock.calls.length, 1);
805+
assert.equal(claudeUsage, null);
806+
assert.equal(routing.claude.readCodexUsage.mock.calls.length, 0);
807+
}),
808+
);
809+
775810
it.effect("routes provider operations and rollback conversation", () =>
776811
Effect.gen(function* () {
777812
const provider = yield* ProviderService;

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ import {
2020
ProviderSessionStartInput,
2121
ProviderStopSessionInput,
2222
type ProviderInstanceId,
23-
type ProviderDriverKind,
2423
type ProviderRuntimeEvent,
2524
type ProviderSession,
25+
ProviderDriverKind,
2626
} from "@t3tools/contracts";
2727
import { Cause, Effect, Layer, Option, PubSub, Ref, Schema, SchemaIssue, Stream } from "effect";
2828

@@ -922,6 +922,24 @@ const makeProviderService = Effect.fn("makeProviderService")(function* (
922922
const getInstanceInfo: ProviderServiceShape["getInstanceInfo"] = (instanceId) =>
923923
registry.getInstanceInfo(instanceId);
924924

925+
const getCodexUsage: ProviderServiceShape["getCodexUsage"] = Effect.fn(
926+
"ProviderService.getCodexUsage",
927+
)(function* (instanceId) {
928+
const info = yield* registry
929+
.getInstanceInfo(instanceId)
930+
.pipe(Effect.catch(() => Effect.succeed(null)));
931+
if (!info || info.driverKind !== ProviderDriverKind.make("codex")) {
932+
return null;
933+
}
934+
const adapter = yield* registry
935+
.getByInstance(instanceId)
936+
.pipe(Effect.catch(() => Effect.succeed(null)));
937+
if (!adapter?.readCodexUsage) {
938+
return null;
939+
}
940+
return yield* adapter.readCodexUsage().pipe(Effect.catch(() => Effect.succeed(null)));
941+
});
942+
925943
const rollbackConversation: ProviderServiceShape["rollbackConversation"] = Effect.fn(
926944
"rollbackConversation",
927945
)(function* (rawInput) {
@@ -1022,6 +1040,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* (
10221040
listSessions,
10231041
getCapabilities,
10241042
getInstanceInfo,
1043+
getCodexUsage,
10251044
rollbackConversation,
10261045
// Each access creates a fresh PubSub subscription so that multiple
10271046
// consumers (ProviderRuntimeIngestion, CheckpointReactor, etc.) each

0 commit comments

Comments
 (0)