feat(rivetkit): c.keepAwake

[NathanFlurry]

No description provided.

[NathanFlurry]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• feat(rivetkit): add clearAlarm #4004
• feat(rivetkit): c.keepAwake #4003 👈 (View in Graphite)
• feat(rivetkit): add run #4002
• feat: actor messages & queues #3989 : 1 other dependent PR (#4005 )
• main

---

How to use the Graphite Merge Queue

Add the label merge-queue to this PR to add it to the merge queue.

You must have a Graphite account in order to use the merge queue. Sign up using this link.

An organization admin has enabled the Graphite Merge Queue in this repository.

Please do not merge from GitHub as this will restart CI on PRs being processed by the merge queue.

This stack of pull requests is managed by Graphite. Learn more about stacking.

[claude]

Pull Request Review: feat(rivetkit): c.keepAwake

Thank you for this contribution! The keepAwake feature is a valuable addition that addresses a real need for preventing actor sleep during critical async operations. Here's my detailed review:

✅ Strengths

1. Clear API design: The method signature and behavior are intuitive and well-documented
2. Proper error handling: The finally block ensures the counter is always decremented
3. Good defensive programming: The negative counter check prevents state corruption
4. Comprehensive documentation: Both code comments and website docs are thorough and include helpful examples
5. Integration with existing sleep logic: Properly integrated into the #canSleep() check at rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts:1110

🔍 Code Quality Issues

1. Race Condition in Counter Management

Location: rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts:721-740

The current implementation has a potential race condition:

async keepAwake<T>(promise: Promise<T>): Promise<T> {
    this.assertReady();

    this.#activeKeepAwakeCount++;  // ← Increment happens immediately
    this.resetSleepTimer();

    try {
        return await promise;  // ← Promise executes
    } finally {
        this.#activeKeepAwakeCount--;  // ← Decrement in finally
        // ...
    }
}

Issue: If a user passes an already-resolved promise, the counter increments and resetSleepTimer() is called, but the promise resolves synchronously before the counter is decremented. This could lead to multiple synchronous keepAwake calls incrementing the counter rapidly and the sleep timer being reset unnecessarily for already-completed work.

Recommendation: The current code is functionally correct since the finally block will execute, but consider documenting this behavior if pre-resolved promises are expected use cases.

2. Inconsistency with waitUntil Design

Comparison:

• waitUntil: Fire-and-forget, errors are caught and logged
• keepAwake: Blocking, errors propagate to caller

Issue: While the different error handling is documented, the naming similarity might confuse users about their relationship. The docs do not clearly explain when to use one vs the other.

Recommendation:

• Add a comparison table in the documentation showing the key differences
• Consider adding a code comment in the implementation explaining the design decision

3. Missing Test Coverage

Critical gap: No tests found for the keepAwake functionality.

Recommended test cases:

1. Verify #activeKeepAwakeCount increments/decrements correctly
2. Verify sleep is prevented while promise is active
3. Verify sleep timer resets after promise completes
4. Verify error propagation (promise rejection)
5. Verify behavior with already-resolved promises
6. Verify behavior with concurrent keepAwake calls
7. Verify counter never goes negative in normal operation
8. Verify integration with #canSleep() check

📝 Documentation Issues

1. Website Documentation

Location: website/src/content/docs/actors/lifecycle.mdx:669-720

Issues:

• The tickActor example at line 686 is good but could be simplified for clarity
• The distinction between keepAwake and waitUntil could be clearer with a comparison table
• Missing guidance on what happens if you nest keepAwake calls

Recommendation: Add a comparison section showing when to use keepAwake vs waitUntil with a table highlighting key differences (use case, error handling, return value, blocking behavior, sleep prevention).

2. JSDoc Comments

Location: rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts:200-207

The JSDoc is good, but could mention:

• What happens if the actor is already stopping when keepAwake is called
• Behavior with nested keepAwake calls
• Performance implications of excessive use

⚠️ Potential Issues

1. No Maximum Limit on Counter

The #activeKeepAwakeCount has no upper bound. While unlikely to be an issue in practice, a runaway loop calling keepAwake could prevent sleep indefinitely.

Recommendation: Consider adding logging when the count exceeds a threshold (e.g., 100) to help debug issues.

2. Sleep Timer Reset Timing

The sleep timer is reset both at the start AND end of the keepAwake call (lines 725 and 738). This is correct behavior, but the double reset might be confusing when debugging.

Recommendation: Add a code comment explaining why the timer is reset twice (once at start to prevent sleep during operation, once at end to restart the idle countdown).

🎯 Recommendations Summary

Must fix before merge:

1. ❌ Add test coverage for the keepAwake functionality

Should fix:
2. ⚠️ Add clarifying comments about the double resetSleepTimer() calls
3. ⚠️ Add comparison table in docs for keepAwake vs waitUntil
4. ⚠️ Add warning log when #activeKeepAwakeCount exceeds a threshold

Nice to have:
5. 💡 Enhance JSDoc with edge case behavior
6. 💡 Add example showing nested keepAwake behavior

🔒 Security Considerations

No security concerns identified. The feature does not introduce any new attack vectors.

⚡ Performance Considerations

The performance impact is minimal:

• Counter increment/decrement: O(1)
• Sleep timer reset: O(1)
• No memory leaks (promise completes and cleanup happens in finally block)

Overall Assessment

This is a well-designed feature that solves a real problem. The main concern is the lack of test coverage, which should be addressed before merging. The implementation is solid, but would benefit from the clarifications mentioned above.

Recommendation: Request changes to add tests, then approve.