chore(rivetkit): close websockets on out of sequence messages #3404
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
3 Skipped Deployments
|
PR Review: Close WebSockets on Out-of-Sequence MessagesSummaryThis PR enhances hibernatable WebSocket reliability by detecting and closing connections when out-of-sequence messages are received (regressed, duplicate, or skipped message indices). The changes are generally well-implemented with good attention to detail. Positive Aspects
Potential Issues & SuggestionsCritical: WebSocket CleanupIssue: When closing the WebSocket with Risk: This could leave stale entries in:
Recommendation: Consider explicitly calling cleanup after closing: // Close the WebSocket and skip processing
wsContext.close(1008, closeReason);
this.#cleanupHibernatableWebSocket(
actorId,
requestIdBuf,
requestId,
"close",
undefined
);
return;Minor: Edge Case - First Message HandlingThe code only checks sequence when Question: Should there be validation that the first message has an expected starting index? DocumentationThe comment at line 617-620 is good, but could be enhanced to mention the new validation behavior: // Handle hibernatable WebSockets:
// - Validate message sequence (close on out-of-sequence)
// - Check for out of sequence messages
// - Save msgIndex for WS restoration
// - Queue WS acksTesting RecommendationsThis change would benefit from driver tests covering:
Consider adding these to Performance ConsiderationThe code now blocks message processing for out-of-sequence messages, which is correct. However, there's no rate limiting on log spam if a buggy client repeatedly sends bad sequences. Consider adding log throttling if this becomes an issue in production. Security Considerations✅ The change improves security by preventing potential replay or sequence manipulation attacks on hibernatable WebSockets. Code QualityThe code follows TypeScript best practices and maintains consistency with the existing codebase style. The structured logging approach aligns with the project's logging patterns. RecommendationApprove with minor changes: Address the cleanup concern and consider adding tests. The core logic is sound and improves system reliability. |
More templates
@rivetkit/actor
@rivetkit/cloudflare-workers
@rivetkit/core
@rivetkit/db
@rivetkit/framework-base
@rivetkit/next-js
@rivetkit/react
rivetkit
@rivetkit/sql-loader
@rivetkit/engine-runner
@rivetkit/engine-runner-protocol
commit: |
NathanFlurry
force-pushed
the
11-06-fix_rivetkit_properly_apply_clientendpoint_from_metadata_lookup
branch
from
e5467bd to
8b4fb49
Compare
NathanFlurry
force-pushed
the
11-06-chore_rivetkit_close_websockets_on_out_of_sequence_messages
branch
from
2368e4d to
0513dc4
Compare
PR Review: Close WebSockets on Out-of-Sequence MessagesSummaryThis PR improves the robustness of hibernatable WebSocket handling by detecting and terminating connections that receive out-of-sequence messages. The changes add proper validation to catch message index regressions, duplicates, and gaps. Positive Aspects1. Improved Error Detection
2. Clear WebSocket Close Codes 3. Proper Control Flow 4. Better Logging Issues & Concerns🔴 Critical: Message Processing Still Occurs After CloseLocation: Issue: After calling Why This Matters:
Recommended Fix: // Option 1: Remove the event listener before closing
websocket.removeEventListener("message", messageHandler);
wsContext.close(1008, closeReason);
return;
// Option 2: Add a flag to prevent further processing
if (this.#closedWebSockets.has(requestId)) {
return;
}You'll need to store a reference to the message handler function to use 🟡 Moderate: First Message HandlingLocation: Issue: When Potential Problems:
Questions to Consider:
Recommended Action: } else {
// First message for this WebSocket
const expectedStartIndex = 0; // or 1, depending on your protocol
if (event.rivetMessageIndex !== expectedStartIndex) {
logger().warn({
msg: "first message has unexpected index",
requestId,
actorId,
receivedIndex: event.rivetMessageIndex,
expectedIndex: expectedStartIndex,
});
// Decide: accept it, reject it, or validate against hibernatableWs.msgIndex
}
this.#hibernatableWebSocketAckQueue.set(requestId, {
requestIdBuf,
messageIndex: event.rivetMessageIndex,
});
}🟡 Moderate: Race Condition in State Update OrderLocation: Observation: The code updates Potential Impact:
Recommendation:
🟢 Minor: WebSocket Close Code 1008Location: Observation: The PR uses close code Note: According to RFC 6455:
The choice of Test Coverage ConcernsMissing Test CasesThe existing test file (
Recommended Tests: test("should close connection on duplicate message index", async (c) => {
// Send message with index N, then send another with index N
// Expect WebSocket to close with code 1008
});
test("should close connection on regressed message index", async (c) => {
// Send message with index N, then send one with index N-1
// Expect WebSocket to close with code 1008
});
test("should close connection on skipped message index", async (c) => {
// Send message with index N, then send one with index N+2
// Expect WebSocket to close with code 1008
});
test("should accept sequential message indices", async (c) => {
// Send messages with indices 0, 1, 2, 3...
// Expect all to be processed normally
});Note: These tests would require either:
Performance ConsiderationsPositive:
Neutral:
Security ConsiderationsPositive:
Potential Concern:
Code QualityStrengths:
Minor Style Notes:
Recommendations SummaryMust Address:
Should Address: Nice to Have: Overall AssessmentThis is a solid improvement to the hibernatable WebSocket reliability. The core logic is sound, but there's a critical issue with the async nature of WebSocket closing that needs to be addressed. Once the message processing race condition is fixed, this will significantly improve protocol robustness. Recommendation: Request changes to address the critical issue, then approve. |
Code Review: WebSocket Out-of-Sequence Message HandlingThank you for this PR! This change adds important safeguards for hibernatable WebSocket message sequencing. Here's my detailed review: ✅ Strengths
🔍 Concerns & Questions1. Potential Race Condition (Medium Priority)// Line 711: Message handler called AFTER validation
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));Issue: If the WebSocket is closed at line 672, there's a potential race where Recommendation: Consider checking if the WebSocket is still open before calling the handler: // After all hibernation logic
if (wsContext.readyState === WebSocket.OPEN) {
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
}2. Change in Behavior for Non-Sequential Messages (High Priority - Breaking Change)Previous behavior: Warned about out-of-sequence messages but still processed them and updated to the highest seen index. New behavior: Immediately closes the connection on any sequence violation. Concerns:
Questions:
3. Missing Test Coverage (High Priority)I couldn't find test cases for the new out-of-sequence detection logic. This is critical functionality that should be tested. Recommendation: Add tests to verify:
Consider adding to 4. Inconsistency in Sequence Validation (Low Priority)The validation only applies when:
Questions:
5. Type Safety (Low Priority)invariant(event.rivetMessageIndex, "missing rivetMessageIndex");
// ...
if (event.rivetMessageIndex !== previousIndex + 1) {The 📊 Performance ConsiderationsThe changes appear performance-neutral:
🔒 Security AssessmentPositive: This change improves security by preventing out-of-sequence message processing, which could potentially be exploited to:
Close code rationale: 1008 is appropriate and follows RFC 6455 guidelines for protocol violations. 📝 Recommendations
🎯 Overall AssessmentThe core logic is sound and improves reliability of hibernatable WebSockets. The main concerns are:
Recommendation: Request tests before merging, and consider adding telemetry to monitor the impact in production. Great work on improving the robustness of WebSocket message handling! Let me know if you'd like me to help write tests for this functionality. |
NathanFlurry
force-pushed
the
11-06-fix_rivetkit_properly_apply_clientendpoint_from_metadata_lookup
branch
from
8b4fb49 to
406bd5a
Compare
NathanFlurry
force-pushed
the
11-06-chore_rivetkit_close_websockets_on_out_of_sequence_messages
branch
from
0513dc4 to
989db0f
Compare
Code ReviewThank you for this PR! This is a critical improvement for handling out-of-sequence WebSocket messages in the hibernatable WebSocket system. Here's my detailed review: ✅ Strengths
🔍 Potential Issues & Suggestions1. Race Condition Risk (Medium Priority)The current code has moved the // Line 675: Close the WebSocket
wsContext.close(1008, closeReason);
return;
// Line 714: This still gets called for valid messages
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));Question: After calling Suggestion: Consider whether you need to track closed connections to prevent processing messages after a close has been initiated. 2. Test Coverage (High Priority)I searched for test coverage and did not find any tests specifically for this hibernatable WebSocket message sequencing logic. This is a critical correctness and security feature that should have comprehensive test coverage. Recommendation: Add tests for:
3. WebSocket Close Code (Low Priority)You're using close code
Question: Is there a documented standard for WebSocket close codes in your system? Consider adding a comment or constant to clarify why 1008 was chosen. 4. Removed Behavior Change (Medium Priority - Needs Verification)The old code had this logic: // Update to the highest seen index
if (event.rivetMessageIndex > previousIndex) {
currentEntry.messageIndex = event.rivetMessageIndex;
}The new code always updates to the received index: // Update to the next index
currentEntry.messageIndex = event.rivetMessageIndex;Question: In the old version, if messages arrived out of order but the connection wasn't closed, it would track the highest index seen. Was this "forgiving" behavior intentional, or was it a bug? The new behavior is stricter (correct for sequential processing), but this is a behavioral change worth confirming. 5. Message Processing After Close (Low Priority - Clarification)At line 714, you're still calling Status: ✅ This appears correct on review, but worth double-checking. 📝 Code Quality
🔒 Security Considerations
🎯 Performance Considerations
📊 SummaryOverall Assessment: This is a solid improvement that enhances the robustness of the hibernatable WebSocket system. The main concern is the lack of test coverage for this critical functionality. Recommended Actions Before Merge:
Optional Improvements:
Great work on improving the reliability of this system! 🚀 |
PR Review: Close WebSockets on Out of Sequence MessagesSummaryThis PR improves the handling of hibernatable WebSocket message sequencing by actively closing connections when out-of-sequence messages are detected, rather than just logging warnings. The change also moves message processing to occur after validation. Positive Changes1. Stronger Sequence Validation ✅The PR changes from a permissive approach (accepting out-of-sequence messages with warnings) to a strict approach (closing connections on sequence violations). This is a good defensive practice for hibernatable WebSockets where message ordering is critical. 2. Clear Error Categorization ✅The code now properly categorizes sequence violations into three types:
3. Improved Control Flow ✅Moving the
4. Better Log Message ✅The warning message now includes "closing connection" which makes the action being taken clear. Issues & Concerns1. Logical Issue: Index Update Should Not Happen on Error
|
Code ReviewSummaryThis PR improves the handling of out-of-sequence WebSocket messages in the hibernatable WebSocket system by immediately closing connections when sequence violations are detected, rather than trying to handle them gracefully. Positive Changes ✅
Issues and Concerns 🔍1. Critical Race Condition (High Priority)The original code called // Line 616: invariants check
invariant(event.rivetRequestId, "missing rivetRequestId");
invariant(event.rivetMessageIndex, "missing rivetMessageIndex");
// Lines 624-711: Hibernation logic that may close the connection
// Line 714: Process message AFTER potential close
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));Problem: If the connection is closed at line 675, the Recommendation: Add a check before calling if (wsContext.readyState !== WebSocket.CLOSED && wsContext.readyState !== WebSocket.CLOSING) {
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
}Or use a flag to track whether the message should be processed. 2. Inconsistent State After Close (Medium Priority)When closing on line 675, the hibernatable WebSocket tracking state is not cleaned up. The entry remains in:
However, the close event handler (line 717) should trigger cleanup via 3. WebSocket Close Code Usage (Low Priority)Using close code
The current 4. Missing First Message Handling (Medium Priority)When Issue: If the first message received has Recommendation: Add validation for the first message: } else {
// First message for this connection
if (event.rivetMessageIndex !== 0) {
logger().warn({
msg: "first hibernatable websocket message has non-zero index",
requestId,
actorId,
firstMessageIndex: event.rivetMessageIndex,
});
}
this.#hibernatableWebSocketAckQueue.set(requestId, {
requestIdBuf,
messageIndex: event.rivetMessageIndex,
});
}5. Timing Concern (Low Priority)The Testing Recommendations 🧪Add tests for:
Example test case locations:
Performance Considerations ⚡No significant performance concerns. The changes:
Security Considerations 🔒Positive: Closing on sequence violations prevents:
Consider: Ensure that legitimate network reordering scenarios don't trigger false positives. WebSocket should maintain message order, but if proxies/load balancers are involved, verify this assumption holds. Style & Conventions ✨The code follows the existing patterns in the file well:
Overall AssessmentThis is a good security improvement but has one critical issue (race condition on line 714) that should be addressed before merging. The fundamental approach of closing on sequence violations is sound and improves system reliability. Recommendation: Request changes to address the race condition, then approve. Suggested Changes
Review completed following CLAUDE.md conventions |
NathanFlurry
force-pushed
the
11-06-chore_rivetkit_close_websockets_on_out_of_sequence_messages
branch
from
aaad3ae to
775d512
Compare
NathanFlurry
force-pushed
the
11-06-fix_rivetkit_properly_apply_clientendpoint_from_metadata_lookup
branch
from
6e38356 to
5270367
Compare
Code Review: Close WebSockets on Out-of-Sequence MessagesSummaryThis PR improves the robustness of hibernatable WebSocket handling by detecting and closing connections when message indices are out of sequence (regressed, duplicate, or skipped). The implementation is solid with minor suggestions for improvement. ✅ Strengths
🔍 Issues & Suggestions1. Race Condition Risk (Medium Priority)The current implementation has a potential race condition. The code now validates and closes BEFORE calling Line 714: If the WebSocket is closed at line 675, the // After line 676, add a flag check or use wsContext.readyState
if (wsContext.readyState !== wsContext.OPEN) {
return;
}
// Or wrap the final onMessage call:
if (wsContext.readyState === wsContext.OPEN) {
wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
}2. Inconsistent State Management (Low Priority)When closing due to out-of-sequence messages, the code doesn't update 3. Missing Test Coverage (High Priority)I searched for tests covering hibernatable WebSocket behavior and found none specifically testing:
Recommendation: Add tests to verify: // Test cases needed:
// 1. Message index 1, 2, 3 -> success
// 2. Message index 1, 2, 1 -> closes with "ws.message_index_regressed"
// 3. Message index 1, 2, 2 -> closes with "ws.message_index_duplicate"
// 4. Message index 1, 2, 5 -> closes with "ws.message_index_skip"
// 5. onMessage should not fire after close on sequence violation4. Close Reason String Format (Low Priority)The close reasons use dot-notation (
🛡️ Security ConsiderationsPositive: This change improves security by preventing potential replay attacks or message reordering attacks. Out-of-sequence messages could indicate:
Closing the connection on sequence violations is the correct defensive approach. ⚡ Performance ConsiderationsNeutral/Positive:
📋 Code QualityMinor: Consider extracting the sequence validation into a helper method to reduce complexity in the event handler: private validateMessageSequence(
requestId: string,
messageIndex: number,
wsContext: any
): boolean {
const currentEntry = this.#hibernatableWebSocketAckQueue.get(requestId);
if (!currentEntry) return true; // First message
const previousIndex = currentEntry.messageIndex;
if (messageIndex === previousIndex + 1) return true; // Valid sequence
// Handle out-of-sequence...
return false;
}🎯 RecommendationsBefore Merging:
Optional: Overall, this is a solid improvement to connection reliability. The main concern is test coverage and the potential race condition with the async |
Code ReviewSummaryThis PR improves the robustness of hibernatable WebSocket message handling by detecting out-of-sequence messages and closing connections when sequencing violations occur. The changes move message processing to happen after validation, preventing potentially corrupted state from being processed. Strengths
Potential Issues and Suggestions1. Race Condition with Early Return (Medium Priority)The early return at line 676 prevents message processing, but wsContext.close() is asynchronous. If close does not complete immediately, subsequent messages might still arrive and be processed. Consider adding a flag to track connections being terminated. 2. Missing Test Coverage (High Priority)The PR lacks tests for out-of-sequence detection logic. Recommended test cases:
3. Invariant Assumptions (Low Priority)Lines 617-618 use invariants on rivetRequestId and rivetMessageIndex, but these fields are marked experimental and optional in websocket-interface.ts. Non-hibernatable WebSockets might not provide these fields. 4. Documentation Gap (Low Priority)The significant behavior change (closing on out-of-sequence) should be documented with comments explaining the rationale. 5. State Management Order (Low Priority)State is updated before message processing. If onMessage throws, state is still advanced. Confirm this is intended behavior. Security ConsiderationsPositive: Closing on sequence violations prevents replay attacks and state corruption. Validation happens before user handlers execute. Performance ConsiderationsMinimal impact: O(1) sequence check only for hibernatable WebSockets. RecommendationConditional Approval - The core fix is solid, but recommend:
Review generated by Claude Code |
Merge activity
|
graphite-app
bot
deleted the
11-06-chore_rivetkit_close_websockets_on_out_of_sequence_messages
branch
No description provided.