diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index c6d1ab4654110..cf31e7b53132b 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -4617,6 +4617,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio snoozedUntil: null, snoozedAt: null, lastSnoozeWakeAt: alreadyAwake ? thread.lastSnoozeWakeAt : now, + limitRecovery: thread.limitRecovery ? { ...thread.limitRecovery, snooze: false } : null, updatedAt: alreadyAwake ? thread.updatedAt : now, }; } @@ -4644,6 +4645,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio snoozedUntil: null, snoozedAt: null, lastSnoozeWakeAt: thread.snoozedUntil == null ? thread.lastSnoozeWakeAt : now, + limitRecovery: thread.limitRecovery ? { ...thread.limitRecovery, snooze: false } : null, updatedAt: alreadyPinned && !promotes ? thread.updatedAt : now, }; } diff --git a/apps/server/src/orchestration-v2/runtimeLayer.test.ts b/apps/server/src/orchestration-v2/runtimeLayer.test.ts index 9e054e3a0bf89..50626f396e09d 100644 --- a/apps/server/src/orchestration-v2/runtimeLayer.test.ts +++ b/apps/server/src/orchestration-v2/runtimeLayer.test.ts @@ -10870,6 +10870,8 @@ it.layer(layerTest)("usage-limit recovery", (it) => { "independent-patches", "expired-snooze", "wake", + "wake-then-resume", + "pin-then-resume", "stop", "stop-before-arm", "stop-before-arm-mismatched-recovery", @@ -11019,8 +11021,10 @@ it.layer(layerTest)("usage-limit recovery", (it) => { assert.isNull(current.thread.snoozedUntil); return; } - const snooze = scenario === "wake" || scenario === "cancel-resume-keep-snooze"; - const autoResume = scenario !== "wake"; + const wakesBeforeResume = scenario === "wake-then-resume" || scenario === "pin-then-resume"; + const snooze = + scenario === "wake" || scenario === "cancel-resume-keep-snooze" || wakesBeforeResume; + const autoResume = scenario !== "wake" && !wakesBeforeResume; const arm = limitRecoveryCommand(shell, autoResume, DateTime.toEpochMillis(now), snooze); assert(arm?.type === "thread.metadata.update" && arm.limitRecovery != null); const automaticArm = { @@ -11193,6 +11197,33 @@ it.layer(layerTest)("usage-limit recovery", (it) => { }); assert.isNull((yield* orchestrator.getThreadProjection(threadId)).thread.snoozedUntil); } + if (wakesBeforeResume) { + yield* orchestrator.dispatch( + scenario === "pin-then-resume" + ? { + type: "thread.pin", + commandId: CommandId.make(`recovery:pin:${scenario}`), + threadId, + } + : { + type: "thread.unsnooze", + commandId: CommandId.make(`recovery:sidebar-wake:${scenario}`), + threadId, + reason: "user", + }, + ); + assert.isNull((yield* orchestrator.getThreadProjection(threadId)).thread.snoozedUntil); + // "Resume at reset" omits snooze, so it inherits the stored flag. + yield* orchestrator.dispatch({ + type: "thread.metadata.update", + commandId: CommandId.make(`recovery:resume-after-wake:${scenario}`), + threadId, + limitRecovery: { runId: run.id, resetAt, autoResume: true }, + }); + const current = yield* orchestrator.getThreadProjection(threadId); + assert.isTrue(current.thread.limitRecovery!.autoResume); + assert.isNull(current.thread.snoozedUntil); + } assert.isNull(limitRecoveryCommand(armedShell, true, DateTime.toEpochMillis(now))); yield* orchestrator.dispatch({ type: "message.dispatch",