diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts index 84c78b91453c2..ff82020155c75 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts @@ -1246,6 +1246,9 @@ describe("CodexAdapterV2 post-settle continuation", () => { transcript.scenario, ).pipe(Effect.orDie); const changes = yield* makeChangeSignal; + const usageLimitsUpdates: Array< + Parameters>[0] + > = []; const continuationRequests: Array = []; const driver = yield* CodexReplay.makeReplayDriver( @@ -1295,11 +1298,14 @@ describe("CodexAdapterV2 post-settle continuation", () => { settings: DEFAULT_CODEX_SETTINGS, environment: {}, clientFactory, - ...(onUsageLimits ? { onUsageLimits } : {}), digest: yield* Digest.Digest, fileSystem, idAllocator, serverConfig, + onUsageLimits: (update) => + Effect.sync(() => { + usageLimitsUpdates.push(update); + }).pipe(Effect.andThen(onUsageLimits?.(update) ?? Effect.void)), continuationRequests: { offer: (request) => Effect.sync(() => { @@ -1363,6 +1369,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { threadId, serverConfig, events, + usageLimitsUpdates, continuationRequests, terminalEvents, subagentUpdates, @@ -5572,6 +5579,54 @@ describe("CodexAdapterV2 post-settle continuation", () => { notification: false, expectedClass: "usage_limit", }, + { + name: "stale-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "empty-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "failed-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "timed-out-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "unexhausted-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "newer-notification-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "omitted-window-read", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, + { + name: "known-reset", + code: "usageLimitExceeded", + notification: false, + expectedClass: "usage_limit", + }, { name: "replacement", code: "usageLimitExceeded", @@ -5707,7 +5762,66 @@ describe("CodexAdapterV2 post-settle continuation", () => { }, ] : []), - ...(scenario.name === "deferred-reset" ? [snapshot] : []), + ...(scenario.name === "deferred-reset" || scenario.name === "known-reset" + ? [snapshot] + : scenario.name === "stale-read" || + scenario.name === "failed-read" || + scenario.name === "unexhausted-read" || + scenario.name === "newer-notification-read" + ? [ + { + ...snapshot, + frame: { + method: "account/rateLimits/updated", + params: { + rateLimits: { + limitId: "codex", + primary: { + usedPercent: 99, + resetsAt: 2000000000, + windowDurationMins: 300, + }, + }, + }, + }, + }, + ] + : []), + ...(scenario.name === "omitted-window-read" + ? [ + { + ...snapshot, + frame: { + method: "account/rateLimits/updated", + params: { + rateLimits: { + limitId: "codex", + primary: { usedPercent: 100, resetsAt: 2000000000 }, + secondary: { usedPercent: 100 }, + }, + }, + }, + }, + ] + : []), + ...(scenario.name === "stale-read" || scenario.name === "empty-read" + ? [ + { + type: "emit_inbound" as const, + label: "account/rateLimits/updated/premium", + frame: { + method: "account/rateLimits/updated", + params: { + rateLimits: { + limitId: "premium", + primary: null, + secondary: null, + }, + }, + }, + }, + ] + : []), ...(scenario.name === "deferred-reset" ? [ { @@ -5791,6 +5905,89 @@ describe("CodexAdapterV2 post-settle continuation", () => { }, }, }, + ...(scenario.expectedClass === "usage_limit" && + !("expectedCode" in scenario) && + scenario.name !== "deferred-reset" && + scenario.name !== "known-reset" && + scenario.name !== "timed-out-read" + ? [ + { + type: "expect_outbound" as const, + label: "account/rateLimits/read", + frame: { + id: 4, + method: "account/rateLimits/read", + params: null, + }, + }, + ...(scenario.name === "newer-notification-read" + ? [ + { + ...snapshot, + frame: { + method: "account/rateLimits/updated", + params: { + rateLimits: { + limitId: "codex", + primary: { + usedPercent: 50, + resetsAt: 2000200000, + windowDurationMins: 300, + }, + }, + }, + }, + }, + ] + : []), + { + type: "emit_inbound" as const, + label: "account/rateLimits/read", + frame: { + id: 4, + ...(scenario.name === "failed-read" + ? { error: { code: -32603, message: "Usage endpoint unavailable." } } + : { + result: { + rateLimits: { limitId: "premium", primary: null, secondary: null }, + rateLimitsByLimitId: { + codex: + scenario.name === "stale-read" || + scenario.name === "empty-read" || + scenario.name === "newer-notification-read" + ? { + limitId: "codex", + primary: { + usedPercent: 100, + resetsAt: 2000000000, + windowDurationMins: 300, + }, + secondary: { + usedPercent: 100, + resetsAt: 2000100000, + windowDurationMins: 10080, + }, + } + : scenario.name === "omitted-window-read" + ? { + limitId: "codex", + secondary: { usedPercent: 99, resetsAt: 2000100000 }, + } + : { + limitId: "codex", + primary: { + usedPercent: 99, + resetsAt: 2000000000, + windowDurationMins: 300, + }, + }, + }, + }, + }), + }, + }, + ] + : []), ...(scenario.name === "late-reset" || "resetAfter" in scenario ? [snapshot] : []), ...(scenario.name === "deferred-reset" ? [ @@ -5829,6 +6026,8 @@ describe("CodexAdapterV2 post-settle continuation", () => { ], }); const resetReceipt = yield* Deferred.make(); + const readStarted = yield* Deferred.make(); + const readInterrupted = yield* Deferred.make(); const snapshotReceipt = yield* Deferred.make(); const harness = yield* makeCodexReplayHarness( transcript, @@ -5838,7 +6037,19 @@ describe("CodexAdapterV2 post-settle continuation", () => { event.turnItem.failure.resetAt === resetAt ? Deferred.succeed(resetReceipt, undefined) : Effect.void, - undefined, + (method) => + method === "account/rateLimits/read" + ? Deferred.succeed(readStarted, undefined).pipe( + Effect.andThen( + scenario.name === "timed-out-read" + ? Effect.never.pipe( + Effect.onInterrupt(() => Deferred.succeed(readInterrupted, undefined)), + ) + : Effect.void, + ), + Effect.asVoid, + ) + : Effect.void, undefined, undefined, () => Deferred.succeed(snapshotReceipt, undefined), @@ -5852,13 +6063,45 @@ describe("CodexAdapterV2 post-settle continuation", () => { attemptId: RunAttemptId.make(`attempt-limit-${scenario.name}`), }), ); + if (scenario.name === "timed-out-read") { + yield* Deferred.await(readStarted); + assert.lengthOf(harness.terminalEvents(), 0); + yield* TestClock.adjust("3 seconds"); + yield* Deferred.await(readInterrupted); + } yield* harness.firstTerminal; const terminal = harness.terminalEvents()[0]; assert.equal(terminal?.status, "failed"); if (terminal?.status !== "failed") return; assert.equal(terminal.failure.class, scenario.expectedClass); assert.equal(terminal.threadDisposition, "reusable"); - if (scenario.name === "deferred-reset") assert.equal(terminal.failure.resetAt, resetAt); + if (scenario.name.endsWith("-read")) yield* Deferred.await(readStarted); + if (scenario.name === "timed-out-read") { + assert.isNull(terminal.failure.resetAt); + assert.isFalse( + harness.events.some( + (event) => + event.type === "turn_item.updated" && + event.turnItem.type === "error" && + event.turnItem.failure.resetAt != null, + ), + ); + } + if (scenario.name === "known-reset") assert.isFalse(yield* Deferred.isDone(readStarted)); + if (["deferred-reset", "known-reset", "stale-read", "empty-read"].includes(scenario.name)) { + assert.equal(terminal.failure.resetAt, resetAt); + } + if ( + [ + "failed-read", + "unexhausted-read", + "omitted-window-read", + "newer-notification-read", + ].includes(scenario.name) + ) { + assert.isNull(terminal.failure.resetAt); + assert.equal(terminal.failure.message, message); + } if (scenario.name === "matching-details") assert.equal(terminal.failure.message, "Detailed provider allowance explanation."); if ("resetAfter" in scenario) { @@ -5872,15 +6115,27 @@ describe("CodexAdapterV2 post-settle continuation", () => { ), ); } - if (scenario.name === "late-reset") { + if (scenario.name === "newer-notification-read") { + assert.deepEqual( + harness.usageLimitsUpdates + .at(-1) + ?.windows.map(({ usedPercent, resetsAt }) => ({ usedPercent, resetsAt })), + [{ usedPercent: 50, resetsAt: "2033-05-20T11:06:40.000Z" }], + ); + } + if (scenario.name === "late-reset" && terminal.failure.resetAt !== resetAt) { yield* Deferred.await(resetReceipt); - const item = harness.events.find( - (event) => - event.type === "turn_item.updated" && - event.turnItem.type === "error" && - event.turnItem.failure.resetAt === resetAt, + } + if (scenario.name === "stale-read" || scenario.name === "empty-read") { + assert.deepEqual( + harness.usageLimitsUpdates + .at(-1) + ?.windows.map(({ id, usedPercent, resetsAt }) => ({ id, usedPercent, resetsAt })), + [ + { id: "primary", usedPercent: 100, resetsAt: "2033-05-18T03:33:20.000Z" }, + { id: "secondary", usedPercent: 100, resetsAt: resetAt }, + ], ); - assert.isDefined(item); } if (scenario.name === "retry" || "willRetry" in scenario) assert.equal(terminal.retry?.attempt, 1); diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts index 07c4e4e72eca7..57ef8da64242b 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts @@ -4491,12 +4491,8 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( }).pipe(Effect.orDie), ); - yield* client.handleServerNotification("account/rateLimits/updated", (payload) => - Effect.gen(function* () { - const previous = yield* Ref.getAndUpdate(rateLimitSnapshot, (previous) => - mergeCodexRateLimits(previous, payload.rateLimits), - ); - const snapshot = yield* Ref.get(rateLimitSnapshot); + const backfillUsageLimitResetAt = Effect.fn("CodexAdapterV2.backfillUsageLimitResetAt")( + function* (snapshot: CodexRateLimitSnapshot | undefined) { for (const item of (yield* Ref.get(limitedTurnItems)).values()) { const resetAt = failureResetAt(item.failure, snapshot); if (resetAt === null || item.failure.resetAt != null) continue; @@ -4514,6 +4510,15 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( turnItem: updated, }); } + }, + ); + + yield* client.handleServerNotification("account/rateLimits/updated", (payload) => + Effect.gen(function* () { + const previous = yield* Ref.getAndUpdate(rateLimitSnapshot, (previous) => + mergeCodexRateLimits(previous, payload.rateLimits), + ); + yield* backfillUsageLimitResetAt(yield* Ref.get(rateLimitSnapshot)); const update = codexRateLimitsToUpdate(payload.rateLimits, previous); if (update && adapterOptions.onUsageLimits) { const checkedAt = DateTime.formatIso(yield* DateTime.now); @@ -5871,7 +5876,35 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( }, ); - const emitRootTerminal = Effect.fnUntraced(function* ( + const resolveUsageLimitResetAt = Effect.fn("CodexAdapterV2.resolveUsageLimitResetAt")( + function* () { + const beforeRead = yield* Ref.get(rateLimitSnapshot); + const cached = codexUsageLimitResetAt(beforeRead); + if (cached !== null) return cached; + // A rejected request may follow a stale percentage or an empty premium snapshot. + const response = yield* client + .request("account/rateLimits/read", null) + .pipe(Effect.timeout("3 seconds"), Effect.option); + if (Option.isSome(response)) { + const snapshot = + response.value.rateLimitsByLimitId?.codex ?? response.value.rateLimits; + // Main-allowance notifications replace this reference; keep any update that won the race. + const applied = yield* Ref.modify(rateLimitSnapshot, (current) => + current === beforeRead && (!snapshot.limitId || snapshot.limitId === "codex") + ? ([true, snapshot] as const) + : ([false, current] as const), + ); + const update = applied ? codexRateLimitsToUpdate(snapshot, beforeRead) : undefined; + if (update && adapterOptions.onUsageLimits) { + const checkedAt = DateTime.formatIso(yield* DateTime.now); + yield* adapterOptions.onUsageLimits({ ...update, checkedAt }); + } + } + return codexUsageLimitResetAt(yield* Ref.get(rateLimitSnapshot)); + }, + ); + + const emitResolvedRootTerminal = Effect.fnUntraced(function* ( context: ActiveCodexTurnContext, event: CodexRootTerminalEvent, ) { @@ -5889,13 +5922,6 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( : event; yield* emitProviderEvent(current); const continued = continuedRuns.get(context.providerTurnId) ?? [context]; - if (current.status !== "completed") { - yield* Ref.update( - terminalizedNonCompletedNativeTurns, - (nativeTurns) => - new Set([...nativeTurns, ...continued.map((turn) => turn.nativeTurnId)]), - ); - } for (const turn of continued) { continuedRuns.delete(turn.providerTurnId); } @@ -5923,6 +5949,40 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( } }); + // Failure ingestion stops at the terminal; finish the RPC off the reader before emitting it. + const emitRootTerminal = Effect.fnUntraced(function* ( + context: ActiveCodexTurnContext, + event: CodexRootTerminalEvent, + ) { + if (event.status !== "completed") { + const continued = continuedRuns.get(context.providerTurnId) ?? [context]; + yield* Ref.update( + terminalizedNonCompletedNativeTurns, + (nativeTurns) => + new Set([...nativeTurns, ...continued.map((turn) => turn.nativeTurnId)]), + ); + } + if ( + event.status === "failed" && + event.failure.class === "usage_limit" && + event.failure.resetAt == null && + !classifyCodexManagedError(event.failure.code)?.usageLimit && + codexUsageLimitResetAt(yield* Ref.get(rateLimitSnapshot)) === null + ) { + yield* resolveUsageLimitResetAt().pipe( + Effect.flatMap((resetAt) => + emitResolvedRootTerminal(context, { + ...event, + failure: { ...event.failure, resetAt }, + }), + ), + Effect.forkIn(scope), + ); + return; + } + yield* emitResolvedRootTerminal(context, event); + }); + const emitOrDeferRootTerminal = Effect.fn("CodexAdapterV2.emitOrDeferRootTerminal")( function* (input: { readonly context: ActiveCodexTurnContext;