fix(server): using a banked usage reset continues threads waiting on that limit

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

@@ -417,6 +417,14 @@ export interface OrchestratorV2Shape {
     readonly threadId: ThreadId;
     readonly limitRecovery: OrchestrationV2LimitRecoveryUpdate;
   }) => Effect.Effect<void, OrchestratorV2Error>;
+  /** Marks recoveries waiting on these provider instances as reset early: their snooze ends and continuations start. */
+  readonly recordLimitReset: (input: {
+    readonly providerInstanceIds: ReadonlyArray<ProviderInstanceId>;
+    /** A run that failed later hit a limit the reset did not clear. */
+    readonly failedBefore: DateTime.Utc;
+    /** A recovery that waits until this time or later waits on a limit the reset did not clear. */
+    readonly resetBefore?: DateTime.Utc;
+  }) => Effect.Effect<void, OrchestratorV2Error>;
   readonly searchThreadStream: (
     input: OrchestrationV2SearchThreadInput,
   ) => Stream.Stream<OrchestrationV2SearchThreadResult, OrchestratorV2Error>;
@@ -649,6 +657,7 @@ function commandThreadId(command: OrchestrationV2ServerCommand): ThreadId {
     case "thread.pull-request-watch.checks-reopen":
     case "thread.pull-request.sync":
     case "thread.title.regeneration.complete":
+    case "thread.limit-recovery.reset":
     case "thread.project.adopt-folder":
     case "thread.runtime-mode.set":
     case "thread.interaction-mode.set":
@@ -5489,6 +5498,73 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
     },
   );
 
+  const dispatchLimitRecoveryReset = Effect.fn("orchestrationV2.dispatch.limitRecoveryReset")(
+    function* (
+      command: Extract<
+        OrchestrationV2ServerCommand,
+        { readonly type: "thread.limit-recovery.reset" }
+      >,
+      events: Ref.Ref<Array<OrchestrationV2DomainEvent>>,
+    ) {
+      const projection = yield* projectionStore
+        .getThreadRecords(command.threadId, ["runs", "turnItems"], { turnItemTypes: ["error"] })
+        .pipe(
+          Effect.mapError(
+            (cause) => new OrchestratorProjectionError({ threadId: command.threadId, cause }),
+          ),
+        );
+      const thread = projection.thread;
+      const recovery = thread.limitRecovery;
+      const run = usageLimitBlockedRun(projection.runs, projection.turnItems, null);
+      const now = yield* DateTime.now;
+      if (
+        recovery?.runId !== command.runId ||
+        recovery.resetAt !== command.resetAt ||
+        recovery.requestId !== command.requestId ||
+        recovery.resetEarly === true ||
+        !(recovery.autoResume || recovery.snooze) ||
+        Date.parse(recovery.resetAt) <= DateTime.toEpochMillis(now) ||
+        run?.id !== recovery.runId ||
+        latestRootProviderFailure(run, projection.turnItems)?.resetAt !== recovery.resetAt ||
+        thread.archivedAt !== null ||
+        thread.deletedAt !== null ||
+        thread.settledOverride === "settled"
+      ) {
+        return yield* new OrchestratorDispatchError({
+          commandId: command.commandId,
+          commandType: command.type,
+          cause: `Thread ${command.threadId} no longer waits on that limit.`,
+        });
+      }
+      yield* emit(
+        events,
+        command,
+      )({
+        type: "thread.metadata-updated",
+        threadId: command.threadId,
+        providerInstanceId: thread.providerInstanceId,
+        occurredAt: now,
+        payload: {
+          ...thread,
+          limitRecovery: {
+            ...recovery,
+            snooze: false,
+            // Without a scheduled continuation the reset only ends the snooze.
+            ...(recovery.autoResume ? { resetEarly: true } : {}),
+            requestId: command.commandId,
+          },
+          // Ends only the snooze this recovery set; one the user chose stays.
+          ...(recovery.snooze &&
+          thread.snoozedUntil != null &&
+          DateTime.toEpochMillis(thread.snoozedUntil) === Date.parse(recovery.resetAt)
+            ? { snoozedUntil: null, snoozedAt: null, lastSnoozeWakeAt: now }
+            : {}),
+          updatedAt: now,
+        },
+      });
+    },
+  );
+
   const dispatchThreadProjectAdoptFolder = Effect.fn(
     "orchestrationV2.dispatch.threadProjectAdoptFolder",
   )(function* (
@@ -6787,7 +6863,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
         recovery.requestId !== command.usageLimitRecoveryRequestId ||
         recovery.runId !== run.id ||
         recovery.resetAt !== failure.resetAt ||
-        Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now) ||
+        (recovery.resetEarly !== true &&
+          Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now)) ||
         projection.thread.archivedAt !== null ||
         projection.thread.deletedAt !== null ||
         projection.thread.settledOverride === "settled" ||
@@ -15033,6 +15110,9 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
       case "provider-session.detach":
         yield* dispatchProviderSessionDetach(command, events, effects);
         break;
+      case "thread.limit-recovery.reset":
+        yield* dispatchLimitRecoveryReset(command, events);
+        break;
       case "thread.project.adopt-folder":
         yield* dispatchThreadProjectAdoptFolder(command, events);
         break;
@@ -16124,6 +16204,34 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
         }),
       );
     }),
+    recordLimitReset: Effect.fn("OrchestratorV2.recordLimitReset")(function* (input) {
+      const waiting = yield* projectionStore
+        .getLimitRecoveriesAwaitingReset({ ...input, now: yield* DateTime.now })
+        .pipe(Effect.mapError(shellProjectionError));
+      // A refusal is stored under its command id, so each reset tries under its own.
+      const reset = DateTime.toEpochMillis(input.failedBefore);
+      yield* Effect.forEach(
+        waiting,
+        ({ threadId, limitRecovery: { runId, resetAt, requestId } }) =>
+          dispatchWithReceipt({
+            type: "thread.limit-recovery.reset",
+            commandId: CommandId.make(`limit-reset:${threadId}:${runId}:${reset}`),
+            threadId,
+            runId,
+            resetAt,
+            ...(requestId === undefined ? {} : { requestId }),
+          }).pipe(
+            // A thread that moved on since the read is refused and keeps its own schedule.
+            Effect.catchCause((cause) =>
+              Effect.logWarning("orchestration-v2.limit-recovery.reset-failed", {
+                threadId,
+                cause,
+              }),
+            ),
+          ),
+        { discard: true },
+      );
+    }),
     searchThreadStream: (input) =>
       projectionStore
         .searchThreadStream(input)

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

@@ -8295,9 +8295,44 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => {
         now: reset,
       })).find((row) => row.id === threadId)!;
       assert.deepEqual(due.limitRecovery, recovery);
+      const awaitingReset = store
+        .getLimitRecoveriesAwaitingReset({
+          now,
+          providerInstanceIds: [original.providerInstanceId],
+          failedBefore: now,
+        })
+        .pipe(Effect.map((rows) => rows.filter((row) => row.threadId === threadId)));
+      assert.deepEqual(yield* awaitingReset, [{ threadId, limitRecovery: recovery }]);
+      const [armedRow] = yield* sql<{
+        payload_json: string;
+      }>`SELECT payload_json FROM orchestration_v2_projection_threads WHERE thread_id = ${threadId}`;
+      // Nothing waits on a reset once it is recorded, cancelled, archived, settled or deleted.
+      for (const change of [
+        sql`payload_json = json_set(payload_json, '$.limitRecovery.resetEarly', json('true'))`,
+        sql`payload_json = json_set(payload_json, '$.limitRecovery.autoResume', json('false'))`,
+        sql`payload_json = json_set(payload_json, '$.archivedAt', ${DateTime.formatIso(now)})`,
+        sql`payload_json = json_set(payload_json, '$.settledOverride', 'settled')`,
+        sql`deleted_at = ${DateTime.formatIso(now)}`,
+      ]) {
+        yield* sql`UPDATE orchestration_v2_projection_threads SET ${change} WHERE thread_id = ${threadId}`;
+        assert.deepEqual(yield* awaitingReset, []);
+        yield* sql`UPDATE orchestration_v2_projection_threads
+          SET payload_json = ${armedRow!.payload_json}, deleted_at = NULL WHERE thread_id = ${threadId}`;
+      }
+      yield* sql`UPDATE orchestration_v2_projection_threads
+        SET payload_json = json_set(payload_json, '$.limitRecovery.resetEarly', json('true'))
+        WHERE thread_id = ${threadId}`;
+      const dueEarly = store
+        .getLimitRecoveryCandidates({ ...recoveryOptions, autoResume: true })
+        .pipe(Effect.map((rows) => rows.find((row) => row.id === threadId)));
+      assert.isDefined(yield* dueEarly);
       yield* sql`INSERT INTO orchestration_v2_projection_runtime_requests
         (runtime_request_id, thread_id, node_id, kind, status, created_at, payload_json)
         VALUES ('limit-shell:pending-request', ${threadId}, ${original.rootNodeId}, 'approval', 'pending', ${DateTime.formatIso(now)}, '{}')`;
+      // A pending approval or question holds back a recovery that was reset early too.
+      assert.isUndefined(yield* dueEarly);
+      yield* sql`UPDATE orchestration_v2_projection_threads
+        SET payload_json = ${armedRow!.payload_json} WHERE thread_id = ${threadId}`;
       assert.isUndefined(
         (yield* store.getLimitRecoveryCandidates({ ...recoveryOptions, now: reset })).find(
           (row) => row.id === threadId,

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

@@ -30,6 +30,7 @@ import type {
   OrchestrationV2ConversationMessage,
   OrchestrationV2ContextTransfer,
   OrchestrationV2DomainEvent,
+  OrchestrationV2LimitRecovery,
   OrchestrationV2OlderSettledCount,
   OrchestrationV2OlderSettledCursor,
   OrchestrationV2ProjectedTurnItem,
@@ -196,6 +197,11 @@ export type ProjectionLimitRecoveryCandidate = Pick<
   | "snoozedUntil"
 >;
 
+export interface ProjectionAwaitedLimitRecovery {
+  readonly threadId: ThreadId;
+  readonly limitRecovery: OrchestrationV2LimitRecovery;
+}
+
 export type ProjectionThreadPullRequests = Pick<
   OrchestrationV2AppThread,
   "id" | "projectId" | "lineage" | "settledOverride" | "settledAt" | "pullRequests"
@@ -622,6 +628,13 @@ export interface ProjectionStoreV2Shape {
   readonly getLimitRecoveryWakeAt: (
     now: DateTime.Utc,
   ) => Effect.Effect<number | null, ProjectionStoreV2Error>;
+  /** Auto-resume or snooze recoveries not yet reset, for runs these instances failed by `failedBefore`. */
+  readonly getLimitRecoveriesAwaitingReset: (options: {
+    readonly now: DateTime.Utc;
+    readonly providerInstanceIds: ReadonlyArray<ProviderInstanceId>;
+    readonly failedBefore: DateTime.Utc;
+    readonly resetBefore?: DateTime.Utc;
+  }) => Effect.Effect<ReadonlyArray<ProjectionAwaitedLimitRecovery>, ProjectionStoreV2Error>;
   readonly getSettlementCandidates: (
     threadId?: ThreadId,
   ) => Effect.Effect<ReadonlyArray<ProjectionSettlementCandidate>, ProjectionStoreV2Error>;
@@ -4636,7 +4649,10 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
                 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.autoResume') = 1
-                AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) <= julianday(${DateTime.formatIso(options.now)})
+                AND (
+                  json_extract(t.payload_json, '$.limitRecovery.resetEarly') = 1
+                  OR 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)})
@@ -4728,6 +4744,48 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
       Effect.mapError((cause) => new ProjectionStoreSetupError({ cause })),
     );
 
+    const getLimitRecoveriesAwaitingReset = Effect.fn(
+      "ProjectionStore.getLimitRecoveriesAwaitingReset",
+    )(
+      function* (
+        options: Parameters<ProjectionStoreV2Shape["getLimitRecoveriesAwaitingReset"]>[0],
+      ) {
+        if (options.providerInstanceIds.length === 0) return [];
+        const resetBefore = options.resetBefore ? DateTime.formatIso(options.resetBefore) : null;
+        const rows = yield* sql<{ readonly payload_json: string }>`
+          SELECT t.payload_json
+          FROM orchestration_v2_projection_threads t
+          INNER JOIN orchestration_v2_projection_runs r
+            ON r.run_id = json_extract(t.payload_json, '$.limitRecovery.runId')
+          WHERE t.deleted_at IS NULL
+            AND json_extract(t.payload_json, '$.archivedAt') IS NULL
+            AND json_extract(t.payload_json, '$.settledOverride') IS NOT 'settled'
+            AND r.status = 'failed'
+            AND r.provider_instance_id IN ${sql.in(options.providerInstanceIds)}
+            AND julianday(COALESCE(r.completed_at, json_extract(t.payload_json, '$.updatedAt'))) <= julianday(${DateTime.formatIso(options.failedBefore)})
+            AND julianday(json_extract(t.payload_json, '$.limitRecovery.resetAt')) > julianday(${DateTime.formatIso(options.now)})
+            AND (
+              ${resetBefore} IS NULL
+              OR julianday(json_extract(t.payload_json, '$.limitRecovery.resetAt')) < julianday(${resetBefore})
+            )
+            AND json_extract(t.payload_json, '$.limitRecovery.resetEarly') IS NOT 1
+            AND (
+              json_extract(t.payload_json, '$.limitRecovery.autoResume') = 1
+              OR json_extract(t.payload_json, '$.limitRecovery.snooze') = 1
+            )
+          ORDER BY t.thread_id
+        `;
+        const waiting: Array<ProjectionAwaitedLimitRecovery> = [];
+        for (const row of rows) {
+          const thread = yield* decodeThreadPayload(row.payload_json);
+          if (thread.limitRecovery != null)
+            waiting.push({ threadId: thread.id, limitRecovery: thread.limitRecovery });
+        }
+        return waiting;
+      },
+      Effect.mapError((cause) => new ProjectionStoreSetupError({ cause })),
+    );
+
     const getRecoveryThreadIds = Effect.fn("ProjectionStore.getRecoveryThreadIds")(
       function* (
         kind: ProjectionRecoveryKind,
@@ -9128,6 +9186,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
       getProviderControlContext: readOnReader(getProviderControlContext),
       getLimitRecoveryCandidates: readOnReader(getLimitRecoveryCandidates),
       getLimitRecoveryWakeAt: readOnReader(getLimitRecoveryWakeAt),
+      getLimitRecoveriesAwaitingReset: readOnReader(getLimitRecoveriesAwaitingReset),
       getRecoveryThreadIds: readOnReader(getRecoveryThreadIds),
       getUnreadableThreadIds: readOnReader(getUnreadableThreadIds),
       getThreadSnapshot: readOnReader(getThreadSnapshot),

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

@@ -212,6 +212,8 @@ it.effect(
           apply: () => Effect.void,
           getLimitRecoveryCandidates: () => Effect.die("unused getLimitRecoveryCandidates"),
           getLimitRecoveryWakeAt: () => Effect.die("unused getLimitRecoveryWakeAt"),
+          getLimitRecoveriesAwaitingReset: () =>
+            Effect.die("unused getLimitRecoveriesAwaitingReset"),
           getOpenRemoteDelegatedTasks: Effect.die("unused getOpenRemoteDelegatedTasks"),
           getAgentAwarenessSnapshot: () => Effect.die("unused awareness snapshot"),
           getShellSnapshot: () => Effect.die("unused getShellSnapshot"),

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

@@ -56,7 +56,7 @@ export function limitRecoveryCommand(
   }
   if (
     !recovery.autoResume ||
-    resetMs > nowMs ||
+    (recovery.resetEarly !== true && resetMs > nowMs) ||
     (thread.snoozedUntil != null && DateTime.toEpochMillis(thread.snoozedUntil) > nowMs)
   )
     return null;