fix(server): limited threads auto-resume when allowance returns early

apps/server/src/orchestration-v2/Orchestrator.ts

@@ -30,6 +30,7 @@ import {
   type ModelSelection,
   OrchestrationV2Command,
   type OrchestrationV2InternalCommand,
+  type OrchestrationV2LimitRecovery,
   type OrchestrationV2LimitRecoveryUpdate,
   type OrchestrationV2ServerCommand,
   type ThreadPullRequestLink,
@@ -759,6 +760,20 @@ function usageLimitHoldsQueue(
   );
 }
 
+// A run that hit the limit again before the previous reset outlasted what the
+// usage reading showed, so its recovery waits for the reset instead of resuming
+// early on that reading once more.
+function limitRecoveryWaitsForReset(
+  previous: OrchestrationV2LimitRecovery | null | undefined,
+  runId: RunId,
+  now: DateTime.Utc,
+): boolean {
+  if (previous == null) return false;
+  return previous.runId === runId
+    ? previous.waitForReset === true
+    : Date.parse(previous.resetAt) > DateTime.toEpochMillis(now);
+}
+
 function latestStableRun(
   projection: Pick<OrchestrationV2ThreadProjection, "runs">,
 ): OrchestrationV2Run | null {
@@ -3159,7 +3174,10 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
         thread.settledOverride === "settled" ||
         run?.id !== command.limitRecovery.runId ||
         failure?.class !== "usage_limit" ||
-        failure.resetAt !== command.limitRecovery.resetAt ||
+        (failure.resetAt ??
+          (thread.limitRecovery?.runId === run.id && thread.limitRecovery.waitForReset === true
+            ? thread.limitRecovery.resetAt
+            : null)) !== command.limitRecovery.resetAt ||
         resetMs <= DateTime.toEpochMillis(run.completedAt ?? run.requestedAt) ||
         projection.runtimeRequests.some((request) => request.status === "pending")
       ) {
@@ -3363,6 +3381,13 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
                       command.limitRecovery.autoResume ?? previousRecovery?.autoResume ?? false,
                     snooze: command.limitRecovery.snooze ?? previousRecovery?.snooze ?? false,
                     requestId: command.commandId,
+                    ...(limitRecoveryWaitsForReset(
+                      thread.limitRecovery,
+                      command.limitRecovery.runId,
+                      now,
+                    )
+                      ? { waitForReset: true }
+                      : {}),
                   };
           return {
             ...thread,
@@ -5018,8 +5043,13 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
         !recovery?.autoResume ||
         recovery.requestId !== command.usageLimitRecoveryRequestId ||
         recovery.runId !== run.id ||
-        recovery.resetAt !== failure.resetAt ||
-        Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now) ||
+        recovery.resetAt !==
+          (failure.resetAt ?? (recovery.waitForReset === true ? recovery.resetAt : null)) ||
+        (Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now) &&
+          (recovery.waitForReset === true ||
+            command.usageLimitClearedAt === undefined ||
+            Date.parse(command.usageLimitClearedAt) <=
+              DateTime.toEpochMillis(run.completedAt ?? run.requestedAt))) ||
         projection.thread.archivedAt !== null ||
         projection.thread.deletedAt !== null ||
         projection.thread.settledOverride === "settled" ||
@@ -5040,14 +5070,38 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
             ...(recovery?.autoResume &&
             recovery.requestId === command.usageLimitRecoveryRequestId &&
             recovery.runId === command.usageLimitContinuationOfRunId &&
-            projection.thread.snoozedUntil != null &&
-            DateTime.toEpochMillis(projection.thread.snoozedUntil) > DateTime.toEpochMillis(now)
+            (projection.runtimeRequests.some((request) => request.status === "pending") ||
+              (projection.thread.snoozedUntil != null &&
+                DateTime.toEpochMillis(projection.thread.snoozedUntil) >
+                  DateTime.toEpochMillis(now)))
               ? { limitRecovery: { ...recovery, requestId: command.commandId } }
               : {}),
           },
         });
         return;
       }
+      if (Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now)) {
+        const event = yield* emit(
+          events,
+          command,
+        )({
+          type: "thread.metadata-updated",
+          threadId: command.threadId,
+          occurredAt: now,
+          payload: {
+            ...projection.thread,
+            limitRecovery: {
+              ...recovery,
+              runId: idAllocator.derive.run({
+                threadId: command.threadId,
+                ordinal: nextRunOrdinal(projection),
+              }),
+              waitForReset: true,
+            },
+          },
+        });
+        projection = applyToProjection(projection, event);
+      }
     }
 
     if (command.restartContinuationOfRunId !== undefined) {
@@ -9613,6 +9667,9 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
           resetAt: failure.resetAt,
           autoResume: false,
           snooze: false,
+          ...(limitRecoveryWaitsForReset(recovery, latest.id, input.now)
+            ? { waitForReset: true }
+            : {}),
         };
       }
     }

apps/server/src/orchestration-v2/ProjectionStore.test.ts

@@ -911,6 +911,101 @@ it.effect("SQL projection keeps background thread updates out of updatedAt", ()
   backgroundThreadUpdatesKeepTimestamp.pipe(Effect.provide(layerTest)),
 );
 
+it.effect.each([
+  ["memory", ProjectionStore.layerMemory],
+  ["SQL", layerTest],
+] as const)("%s limit recovery loads armed threads before their reset", ([, layer]) =>
+  Effect.gen(function* () {
+    const store = yield* ProjectionStore.ProjectionStoreV2;
+    const threadId = yield* addRolledBackRecoveryCandidate("armed-before-reset");
+    const projection = yield* store.getThreadProjection(threadId);
+    const run = projection.runs[0]!;
+    const now = yield* DateTime.now;
+    const resetAt = "2099-01-01T00:00:00.000Z";
+    yield* store.apply({
+      id: EventId.make("event:armed-before-reset:error"),
+      type: "turn-item.updated",
+      threadId,
+      occurredAt: now,
+      payload: {
+        id: TurnItemId.make("armed-before-reset:error"),
+        threadId,
+        runId: run.id,
+        nodeId: run.rootNodeId,
+        providerThreadId: null,
+        providerTurnId: null,
+        nativeItemRef: null,
+        parentItemId: null,
+        ordinal: 2,
+        status: "failed",
+        title: "Usage limit reached",
+        startedAt: now,
+        completedAt: now,
+        updatedAt: now,
+        type: "error",
+        failure: {
+          class: "usage_limit",
+          message: "Plan limit reached.",
+          resetAt,
+          code: "usageLimitExceeded",
+          retryable: null,
+        },
+      },
+    });
+    yield* store.apply({
+      id: EventId.make("event:armed-before-reset:failed"),
+      type: "run.updated",
+      threadId,
+      occurredAt: now,
+      payload: { ...run, status: "failed" },
+    });
+    const recovery = {
+      runId: run.id,
+      resetAt,
+      autoResume: true,
+      requestId: CommandId.make("recovery:armed-before-reset"),
+    };
+    const updateThread = (suffix: string, patch: Partial<typeof projection.thread>) =>
+      store.apply({
+        id: EventId.make(`event:armed-before-reset:${suffix}`),
+        type: "thread.metadata-updated",
+        threadId,
+        occurredAt: now,
+        payload: { ...projection.thread, limitRecovery: recovery, ...patch },
+      });
+    const candidate = Effect.map(
+      store.getLimitRecoveryCandidates({ now, autoResume: false, snooze: false }),
+      (rows) => rows.find((row) => row.id === threadId),
+    );
+
+    yield* updateThread("armed", {});
+    assert.deepEqual((yield* candidate)?.limitRecovery, recovery);
+    const retry = { ...recovery, resetAt: "2099-01-02T00:00:00.000Z", waitForReset: true };
+    yield* updateThread("retry", { limitRecovery: retry });
+    assert.deepEqual((yield* candidate)?.limitRecovery, retry);
+    const failure = (yield* store.getThreadProjection(threadId)).turnItems.find(
+      (item) => item.id === TurnItemId.make("armed-before-reset:error"),
+    );
+    assert(failure?.type === "error");
+    yield* store.apply({
+      id: EventId.make("event:armed-before-reset:no-reset"),
+      type: "turn-item.updated",
+      threadId,
+      occurredAt: now,
+      payload: { ...failure, failure: { ...failure.failure!, resetAt: null } },
+    });
+    assert.equal((yield* candidate)?.usageLimitResetAt, retry.resetAt);
+    yield* updateThread("no-reset-unarmed", {});
+    assert.isUndefined(yield* candidate);
+    yield* updateThread("retry-cancelled", { limitRecovery: { ...retry, autoResume: false } });
+    assert.isUndefined(yield* candidate);
+    yield* updateThread("snoozed", { snoozedUntil: DateTime.makeUnsafe(resetAt) });
+    assert.isUndefined(yield* candidate);
+    yield* updateThread("cancelled", { limitRecovery: { ...recovery, autoResume: false } });
+    assert.isUndefined(yield* candidate);
+  }).pipe(Effect.provide(layer)),
+);
+
 it.effect("memory recovery selection includes unfinished items from missing runs", () =>
   Effect.gen(function* () {
     const projectionStore = yield* ProjectionStore.ProjectionStoreV2;
@@ -3892,6 +3987,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => {
         if (lastErrorClass === "usage_limit") {
           assert.deepEqual(candidate, {
             id: sqlShell.id,
+            providerInstanceId: sqlShell.providerInstanceId,
             status: sqlShell.status,
             lastErrorClass: sqlShell.lastErrorClass,
             usageLimitResetAt: sqlShell.usageLimitResetAt,
@@ -4070,10 +4166,12 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => {
       yield* sql`UPDATE orchestration_v2_projection_threads
         SET payload_json = json_set(payload_json, '$.limitRecovery', json(${encodeUnknownJsonString(recovery)}))
         WHERE thread_id = ${threadId}`;
-      assert.isUndefined(
+      // Armed rows load before the reset too, so the sweep can see an allowance restored early.
+      assert.deepEqual(
         (yield* store.getLimitRecoveryCandidates({ ...recoveryOptions, autoResume: true })).find(
           (row) => row.id === threadId,
-        ),
+        )?.limitRecovery,
+        recovery,
       );
       const due = (yield* store.getLimitRecoveryCandidates({
         ...recoveryOptions,

apps/server/src/orchestration-v2/ProjectionStore.ts

@@ -171,6 +171,7 @@ export type ProjectionRecoveryKind =
 export type ProjectionLimitRecoveryCandidate = Pick<
   OrchestrationV2ThreadShell,
   | "id"
+  | "providerInstanceId"
   | "status"
   | "lastErrorClass"
   | "latestRunId"
@@ -3862,9 +3863,16 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
           readonly completed_at: string | null;
           readonly failure_payload_json: string;
           readonly last_error: string | null;
+          readonly recovery_reset_at: string | null;
         }>`
           SELECT t.payload_json, r.run_id, r.completed_at,
             item.payload_json AS failure_payload_json,
+            COALESCE(json_extract(item.payload_json, '$.failure.resetAt'), CASE
+              WHEN json_extract(t.payload_json, '$.limitRecovery.runId') IS r.run_id
+                AND json_extract(t.payload_json, '$.limitRecovery.autoResume') = 1
+                AND json_extract(t.payload_json, '$.limitRecovery.waitForReset') = 1
+              THEN json_extract(t.payload_json, '$.limitRecovery.resetAt')
+            END) AS recovery_reset_at,
             (
               SELECT json_extract(session.payload_json, '$.lastError')
               -- Keep bindings outermost: one provider instance can hold every session.
@@ -3910,14 +3918,16 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
             AND json_extract(t.payload_json, '$.archivedAt') IS NULL
             AND json_extract(t.payload_json, '$.settledOverride') IS NOT 'settled'
             AND json_extract(item.payload_json, '$.failure.class') = 'usage_limit'
-            AND json_extract(item.payload_json, '$.failure.resetAt') IS NOT NULL
-            AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) > julianday(COALESCE(r.completed_at, json_extract(t.payload_json, '$.updatedAt')))
+            AND recovery_reset_at IS NOT NULL
+            AND julianday(recovery_reset_at) > julianday(COALESCE(r.completed_at, json_extract(t.payload_json, '$.updatedAt')))
             AND (
               (
                 json_extract(t.payload_json, '$.limitRecovery.runId') IS r.run_id
-                AND json_extract(t.payload_json, '$.limitRecovery.resetAt') IS json_extract(item.payload_json, '$.failure.resetAt')
+                AND (
+                  json_extract(t.payload_json, '$.limitRecovery.resetAt') IS recovery_reset_at
+                  OR json_extract(t.payload_json, '$.limitRecovery.waitForReset') = 1
+                )
                 AND json_extract(t.payload_json, '$.limitRecovery.autoResume') = 1
-                AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) <= julianday(${DateTime.formatIso(options.now)})
                 AND (
                   json_extract(t.payload_json, '$.snoozedUntil') IS NULL
                   OR julianday(json_extract(t.payload_json, '$.snoozedUntil')) <= julianday(${DateTime.formatIso(options.now)})
@@ -3926,11 +3936,11 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
               OR (
                 (
                   json_extract(t.payload_json, '$.limitRecovery.runId') IS NOT r.run_id
-                  OR json_extract(t.payload_json, '$.limitRecovery.resetAt') IS NOT json_extract(item.payload_json, '$.failure.resetAt')
+                  OR json_extract(t.payload_json, '$.limitRecovery.resetAt') IS NOT recovery_reset_at
                 )
                 AND (
                   ${booleanInt(options.autoResume)}
-                  OR (${booleanInt(options.snooze)} AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) > julianday(${DateTime.formatIso(options.now)}))
+                  OR (${booleanInt(options.snooze)} AND julianday(recovery_reset_at) > julianday(${DateTime.formatIso(options.now)}))
                 )
               )
             )
@@ -3951,9 +3961,10 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
           if (summary.lastErrorClass !== "usage_limit") continue;
           candidates.push({
             id: thread.id,
+            providerInstanceId: thread.providerInstanceId,
             status: "failed",
             lastErrorClass: summary.lastErrorClass,
-            usageLimitResetAt: summary.usageLimitResetAt,
+            usageLimitResetAt: row.recovery_reset_at,
             latestRunId: RunId.make(row.run_id),
             latestRunCompletedAt:
               row.completed_at === null ? null : DateTime.makeUnsafe(row.completed_at),
@@ -7071,6 +7082,16 @@ export const layerMemory: Layer.Layer<ProjectionStoreV2> = Layer.effect(
                   thread.settledOverride !== "settled",
               )
               .map(threadShellFromProjection)
+              .map((thread) => ({
+                ...thread,
+                usageLimitResetAt:
+                  thread.usageLimitResetAt ??
+                  (thread.limitRecovery?.runId === thread.latestRunId &&
+                  thread.limitRecovery.autoResume &&
+                  thread.limitRecovery.waitForReset === true
+                    ? thread.limitRecovery.resetAt
+                    : null),
+              }))
               .filter(
                 (thread) =>
                   thread.status === "failed" &&
@@ -7088,13 +7109,15 @@ export const layerMemory: Layer.Layer<ProjectionStoreV2> = Layer.effect(
                   return false;
                 if (
                   thread.limitRecovery?.runId !== thread.latestRunId ||
-                  thread.limitRecovery.resetAt !== thread.usageLimitResetAt
+                  (thread.limitRecovery.resetAt !== thread.usageLimitResetAt &&
+                    !(
+                      thread.limitRecovery.autoResume && thread.limitRecovery.waitForReset === true
+                    ))
                 ) {
                   return options.autoResume || (options.snooze && resetMs > nowMs);
                 }
                 return (
                   thread.limitRecovery.autoResume &&
-                  resetMs <= nowMs &&
                   (thread.snoozedUntil == null ||
                     DateTime.toEpochMillis(thread.snoozedUntil) <= nowMs)
                 );